{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64d4fa0c",
   "metadata": {},
   "source": [
    "# 数据处理\n",
    "- 文档用来测试YOLO数据集的读取与增强\n",
    "- dataloader,augment ……"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ac17c79",
   "metadata": {},
   "source": [
    "## 环境配置与验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fbe2820",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 开启笔记本扩展\n",
    "# 最新版本已失效\n",
    "pip install jupyter_contrib_nbextensions\n",
    "pip install jupyter_nbextensions_configurator\n",
    "jupyter contrib nbextension install --user\n",
    "jupyter nbextensions_configurator enable --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6e18f85",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-18T08:15:52.993163Z",
     "start_time": "2023-12-18T08:15:52.863857Z"
    }
   },
   "outputs": [],
   "source": [
    "# 查看系统GPU驱动\n",
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8dcbfe1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.165274Z",
     "start_time": "2022-07-16T08:19:43.165264Z"
    }
   },
   "outputs": [],
   "source": [
    "# 查看pytorch与CUDA状态\n",
    "import torch\n",
    "\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cac3720c",
   "metadata": {},
   "source": [
    "## yolo标签数据格式\n",
    "\n",
    "- 标注含义\n",
    "<!-- ![yolo_mark](./saved_pic/yolo_mark.webp) -->\n",
    "<img src=\"./saved_pic/yolo_mark.webp\" style=\"width:400px\"/>\n",
    "- 标注文件内容\n",
    "<img src=\"./saved_pic/mark.webp\" style=\"width:350px\"/>\n",
    "- Yolo文件格式\n",
    "<img src=\"./saved_pic/yolo_data.webp\" style=\"width:350px\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5daa4695",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-04-25T09:32:06.908182Z",
     "start_time": "2022-04-25T09:32:06.905327Z"
    }
   },
   "source": [
    "- 自制数据集在Yolo工程中的存放方式\n",
    "```ascii\n",
    "  parent\n",
    "     ├── yolov5\n",
    "     └── datasets\n",
    "             └── car \n",
    "                  ├──images\n",
    "                  |    └──train\n",
    "                  └──labels\n",
    "                       └──train \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4747bc1",
   "metadata": {},
   "source": [
    "```ascii\n",
    "/home/wssysh/data/datasets/bs/\n",
    "├── images\n",
    "│   ├── train\n",
    "│   └── valid\n",
    "└── labels\n",
    "    ├── train\n",
    "    └── valid\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd5fed02",
   "metadata": {},
   "source": [
    "- 三种主流数据格式转换[blog](https://blog.csdn.net/qq_43701912/article/details/120244503#31_YOLO_423)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd6b11fb",
   "metadata": {},
   "source": [
    "- 数据分组[blog](https://www.jianshu.com/p/86c6c57ad699)\n",
    "\n",
    "- 文件夹结构建立与数据集划分"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28380201",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.104496Z",
     "start_time": "2022-07-16T08:19:42.761664Z"
    }
   },
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "import os, random, shutil\n",
    "\n",
    "# # root path\n",
    "# src_path = \"/home/wssysh/data/src_git/yolov5/car/\"\n",
    "# target_path = \"/home/wssysh/data/datasets/bs/\"\n",
    "\n",
    "# # src img and label path\n",
    "# img_path = os.path.join(src_path, 'images/train')\n",
    "# label_path = os.path.join(src_path, 'labels/train')\n",
    "# root path\n",
    "src_path = \"..\"\n",
    "target_path = \"D:/to_learn/bs/\"\n",
    "\n",
    "# src img and label path\n",
    "img_path = os.path.join(src_path, 'pics')\n",
    "label_path = os.path.join(src_path, 'labels/')\n",
    "\n",
    "# 目标文件夹路径\n",
    "train_img_path = os.path.join(target_path, 'images/train')\n",
    "train_label_path = os.path.join(target_path, 'labels/train')\n",
    "val_img_path = os.path.join(target_path, \"images/valid\")\n",
    "val_label_path = os.path.join(target_path, \"labels/valid\")\n",
    "test_img_path = os.path.join(target_path, \"images/test\")\n",
    "test_label_path = os.path.join(target_path, \"labels/test\")\n",
    "\n",
    "paths = [\n",
    "    target_path, target_path + \"images\", target_path + \"labels\",\n",
    "    train_img_path, train_label_path, val_img_path, val_label_path,\n",
    "    test_img_path, test_label_path\n",
    "]\n",
    "\n",
    "# 建立目标文件夹\n",
    "for target_dir in paths:\n",
    "    if not os.path.exists(target_dir):\n",
    "        os.mkdir(target_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba22c917",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.115183Z",
     "start_time": "2022-07-16T08:19:43.115172Z"
    }
   },
   "outputs": [],
   "source": [
    "# 生成目录结构\n",
    "# !tree -d /home/wssysh/data/datasets/bs/\n",
    "!tree D://to_learn/bs/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e4ad7ca",
   "metadata": {},
   "source": [
    "- 以下两个Cell为一体执行"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02d2d675",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.116306Z",
     "start_time": "2022-07-16T08:19:43.116294Z"
    }
   },
   "outputs": [],
   "source": [
    "# linux\n",
    "# 是否获取验证集 一定要先获取 train data\n",
    "get_val_data = True\n",
    "\n",
    "# 读取标签文件\n",
    "label_list = os.listdir(label_path)\n",
    "print(\"共获取有效标签数目：{}\".format(len(label_list)))\n",
    "\n",
    "## 该部分复制标记图像数据集到train\n",
    "str1 = \"cd \" + label_path\n",
    "str2 = \"cp * \" + train_label_path\n",
    "sh_str = str1 + \"&&\" + str2\n",
    "print(sh_str)\n",
    "os.system(sh_str)\n",
    "label_list = os.listdir(train_label_path)\n",
    "\n",
    "img_list = []\n",
    "for label in label_list:\n",
    "    img_list.append(label.replace(\"txt\", \"jpg\"))\n",
    "# 调用 linux sh\n",
    "t_img_names = \" \".join(img_list)\n",
    "str1 = \"cd \" + img_path\n",
    "str2 = \"cp \" + t_img_names + \" \" + train_img_path\n",
    "sh_str = str1 + \"&&\" + str2\n",
    "os.system(sh_str)\n",
    "pic_num = len(os.listdir(train_img_path))\n",
    "print(\"共获取有效图像数目：{}\".format(pic_num))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46a5f584",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.117229Z",
     "start_time": "2022-07-16T08:19:43.117218Z"
    }
   },
   "outputs": [],
   "source": [
    "# windows\n",
    "# 是否获取验证集 一定要先获取 train data\n",
    "get_val_data = True\n",
    "\n",
    "# 读取标签文件\n",
    "label_list = os.listdir(label_path)\n",
    "print(\"共获取有效标签数目：{}\".format(len(label_list)))\n",
    "\n",
    "## 该部分复制标记图像数据集到train\n",
    "%cd ../labels/\n",
    "%pwd\n",
    "cmd_str1 = \"copy * \" + train_label_path.replace(\"/\",os.sep)\n",
    "print(cmd_str1)\n",
    "os.system(cmd_str1)\n",
    "# 回到初始目录\n",
    "%cd D://to_learn/notebook_Yolo/\n",
    "    \n",
    "label_list = os.listdir(train_label_path)\n",
    "img_list = []\n",
    "for label in label_list:\n",
    "    img_list.append(label.replace(\"txt\", \"jpg\"))\n",
    "for img in img_list:\n",
    "    shutil.move(os.path.join(img_path,img),train_img_path)\n",
    "pic_num = len(os.listdir(train_img_path))\n",
    "print(\"共获取有效图像数目：{}\".format(pic_num))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32cc3689",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.118075Z",
     "start_time": "2022-07-16T08:19:43.118064Z"
    }
   },
   "outputs": [],
   "source": [
    "## 该部分为拆分数据集得到val\n",
    "if get_val_data:\n",
    "\n",
    "    rate = 0.2  # rate划分比例\n",
    "    picknumber = int(pic_num * rate)  # 划分验证集图片数量（整数）\n",
    "    img_sample = random.sample(img_list, picknumber)  # 在图片名list中随机选取\n",
    "\n",
    "    label_sample = []  # 选取对应的label\n",
    "    for img_name in img_sample:\n",
    "        label_sample.append(img_name.replace(\"jpg\", \"txt\"))\n",
    "    # 验证是否获取一致\n",
    "    if len(img_sample) != len(label_sample):\n",
    "        print(\"error!\")\n",
    "    # linux sh: mv img_names target_path\n",
    "    # move img\n",
    "    img_names = \" \".join(img_sample)\n",
    "    str1 = \"cd \" + train_img_path\n",
    "    str2 = \"mv \" + img_names + \" \" + val_img_path\n",
    "    sh_str = str1 + \"&&\" + str2\n",
    "    os.system(sh_str)\n",
    "    # move label\n",
    "    label_names = \" \".join(label_sample)\n",
    "    str1 = \"cd \" + train_label_path\n",
    "    str2 = \"mv \" + label_names + \" \" + val_label_path\n",
    "    sh_str = str1 + \"&&\" + str2\n",
    "    os.system(sh_str)\n",
    "    val_num = len(os.listdir(val_img_path))\n",
    "    print(\"拆分完毕,共拆分验证图像数目为：{}\".format(val_num))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f5d09d8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.118889Z",
     "start_time": "2022-07-16T08:19:43.118878Z"
    }
   },
   "outputs": [],
   "source": [
    "# win 可采用此种方式，替换linux sh 操作\n",
    "# 另一种循环复制与移动方式\n",
    "###########################################################################\n",
    "#     print(sample)\n",
    "#     for name in sample:\n",
    "#         jpg_name = name + \".jpg\"\n",
    "#         txt_name = name + \".txt\"\n",
    "#         shutil.move(dir_path + \"/\" + jpg_name, image_save_path + jpg_name)\n",
    "#         shutil.move(dir_path + \"/\" + txt_name, label_save_path + txt_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1416ce51",
   "metadata": {},
   "source": [
    "## 图片格式转换\n",
    "- HEIC2jpg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cd965f3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.120636Z",
     "start_time": "2022-07-16T08:19:43.120625Z"
    }
   },
   "outputs": [],
   "source": [
    "%pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afb452a7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-15T08:06:38.014775Z",
     "start_time": "2022-11-15T08:06:21.029238Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# func HEIC2JPG\n",
    "# 需安装库\n",
    "# !pip install whatimage pyheif\n",
    "import os\n",
    "import whatimage\n",
    "import pyheif\n",
    "import traceback\n",
    "from PIL import Image\n",
    "\n",
    "# 解码转换\n",
    "def decodeImage(bytesIo, path):\n",
    "    try:\n",
    "        fmt = whatimage.identify_image(bytesIo)\n",
    "        if fmt in ['heic']:\n",
    "            img_heif = pyheif.read_heif(bytesIo)\n",
    "            img_pil = Image.frombytes(mode=img_heif.mode, size=img_heif.size, data=img_heif.data)\n",
    "            img_pil.save(path, format=\"jpeg\")\n",
    "    except:\n",
    "        traceback.print_exc()\n",
    "\n",
    "# 读取数据\n",
    "def read_image_file_rb(file_path):\n",
    "    with open(file_path, 'rb') as f:\n",
    "        file_data = f.read()\n",
    "    return file_data\n",
    "\n",
    "# 批量转换\n",
    "if __name__ == \"__main__\":\n",
    "    # HEIC图片所在路径\n",
    "    file_path = \"/home/wssysh/Telegram/\"\n",
    "    for file in os.listdir(file_path):\n",
    "        print('file', file)\n",
    "        data = read_image_file_rb(os.path.join(file_path,file))\n",
    "        # 修改输出路径\n",
    "        decodeImage(data, \"/home/wssysh/data/to_learn/notebooks/pics/date_20221115/\"+file.replace(\"HEIC\",\"jpg\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "804ea7d3",
   "metadata": {},
   "source": [
    "## 数据读取"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56811f61",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.122306Z",
     "start_time": "2022-07-16T08:19:43.122296Z"
    }
   },
   "outputs": [],
   "source": [
    "from torch.utils.data.dataset import Dataset\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4e0096d",
   "metadata": {},
   "source": [
    "### Dataset\n",
    "- Dataset是包装类，将数据包装为Dataset类，然后传入DataLoader中;\n",
    "- 用户想要加载自定义的数据时，只需要继承这个类，并且覆写其中的两个方法\n",
    "1. `__len__`：实现`len(dataset)`，返回整个数据集的大小\n",
    "2. `__getitem__`：用来获取一些索引的数据，使`dataset[i]`返回数据集中第i个样本\n",
    "- 不覆写这两个方法会直接返回错误\n",
    "\n",
    "```python\n",
    "class YoloDataset(Dataset):\n",
    "    def __init__(self, annotation_lines, input_shape, num_classes, train):\n",
    "        super(YoloDataset, self).__init__()\n",
    "        ...\n",
    "\n",
    "    def __len__(self):\n",
    "        ...\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        ...\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe5062be",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.123457Z",
     "start_time": "2022-07-16T08:19:43.123446Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "from PIL import ExifTags, Image\n",
    "from torch.utils.data import Dataset\n",
    "from tqdm import tqdm\n",
    "\n",
    "help_url = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'\n",
    "img_formats = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.dng']\n",
    "vid_formats = ['.mov', '.avi', '.mp4', '.mkv']  # not completed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1adec1f6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.124341Z",
     "start_time": "2022-07-16T08:19:43.124331Z"
    }
   },
   "outputs": [],
   "source": [
    "# read yolo data(img/label) for training/testing\n",
    "class LoadImagesAndLabels(Dataset):\n",
    "\n",
    "    def __init__(self,\n",
    "                 path,\n",
    "                 img_size=640,\n",
    "                 batch_size=16,\n",
    "                 augment=False,\n",
    "                 hyp=None,\n",
    "                 rect=False,\n",
    "                 image_weights=False,\n",
    "                 cache_images=False,\n",
    "                 single_cls=False,\n",
    "                 stride=32,\n",
    "                 pad=0.0):\n",
    "        # -------------- Path 健壮性处理 ---------------- #\n",
    "        try:\n",
    "            f = []\n",
    "            # 获取数据集路径path，包含图片路径的txt文件 or 包含图片的文件夹路径\n",
    "            # 使用pathlib.Path生成与操作系统无关的路径，因为不同操作系统路径的‘/’会有所不同\n",
    "            for p in path if isinstance(path, list) else [path]:\n",
    "                p = str(Path(p))  # os-agnostic\n",
    "                # 获取数据集路径的上级父目录，os.sep为路径里的破折号(不同系统路径破折号不同，os.sep根据系统自适应)\n",
    "                parent = str(Path(p).parent) + os.sep\n",
    "\n",
    "                # 如果路径path为包含图片的文件夹路径\n",
    "                if os.path.isdir(p):  # folder\n",
    "                    f += glob.iglob(p + os.sep + '*.*')\n",
    "\n",
    "                # 如果路径path为包含图片路径的txt文件\n",
    "                elif os.path.isfile(p):  # file\n",
    "                    with open(p, 'r') as t:\n",
    "                        # 获取图片路径，更换相对路径\n",
    "                        t = t.read().splitlines()\n",
    "                        f += [\n",
    "                            x.replace('./', parent)\n",
    "                            if x.startswith('./') else x for x in t\n",
    "                        ]  # local to global path\n",
    "                else:\n",
    "                    raise Exception('%s does not exist' % p)\n",
    "            path = p  # *.npy dir npy文件将零碎的标签整合成整个文件\n",
    "            # 破折号替换为os.sep，os.path.splitext(x)将文件名与扩展名分开并返回一个列表\n",
    "            self.img_files = [\n",
    "                x.replace('/', os.sep) for x in f\n",
    "                if os.path.splitext(x)[-1].lower() in img_formats\n",
    "            ]\n",
    "        except Exception as e:\n",
    "            raise Exception('Error loading data from %s: %s\\nSee %s' %\n",
    "                            (path, e, help_url))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f330ef5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.125322Z",
     "start_time": "2022-07-16T08:19:43.125311Z"
    }
   },
   "outputs": [],
   "source": [
    "# 测试Path健壮性\n",
    "path = \"/home/wssysh/data/datasets/bs/images/train/\"\n",
    "try:\n",
    "    f = []\n",
    "    # 获取数据集路径path，包含图片路径的txt文件 or 包含图片的文件夹路径\n",
    "    # 使用pathlib.Path生成与操作系统无关的路径，因为不同操作系统路径的‘/’会有所不同\n",
    "    for p in path if isinstance(path, list) else [path]:\n",
    "        p = str(Path(p))  # os-agnostic\n",
    "        # 获取数据集路径的上级父目录，os.sep为路径里的破折号(不同系统路径破折号不同，os.sep根据系统自适应)\n",
    "        parent = str(Path(p).parent) + os.sep\n",
    "\n",
    "        # 如果路径path为包含图片的文件夹路径\n",
    "        if os.path.isdir(p):  # folder\n",
    "            f += glob.iglob(p + os.sep + '*.*')\n",
    "\n",
    "        # 如果路径path为包含图片路径的txt文件\n",
    "        elif os.path.isfile(p):  # file\n",
    "            with open(p, 'r') as t:\n",
    "                # 获取图片路径，更换相对路径\n",
    "                t = t.read().splitlines()\n",
    "                f += [\n",
    "                    x.replace('./', parent) if x.startswith('./') else x\n",
    "                    for x in t\n",
    "                ]  # local to global path\n",
    "        else:\n",
    "            raise Exception('%s does not exist' % p)\n",
    "    path = p  # *.npy dir npy文件将零碎的标签整合成整个文件\n",
    "    # 破折号替换为os.sep，os.path.splitext(x)将文件名与扩展名分开并返回一个列表\n",
    "    img_files = [\n",
    "        x.replace('/', os.sep) for x in f\n",
    "        if os.path.splitext(x)[-1].lower() in img_formats\n",
    "    ]\n",
    "except Exception as e:\n",
    "    raise Exception('Error loading data from %s: %s\\nSee %s' %\n",
    "                    (path, e, help_url))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73406e8b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.126111Z",
     "start_time": "2022-07-16T08:19:43.126100Z"
    }
   },
   "outputs": [],
   "source": [
    "# print(img_files)\n",
    "# -------------- 参数初始化 ---------------- #\n",
    "batch_size = 32\n",
    "# 数据集的数量\n",
    "img_num = len(img_files)\n",
    "assert img_num > 0, 'No images found in %s. See %s' % (path, help_url)\n",
    "# 获取batch的索引\n",
    "bi = np.floor(np.arange(img_num) / batch_size).astype(np.int)  # batch index\n",
    "# 一个轮次batch的数量\n",
    "nb = bi[-1] + 1  # number of batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87bf9423",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.127253Z",
     "start_time": "2022-07-16T08:19:43.127236Z"
    }
   },
   "outputs": [],
   "source": [
    "print(bi)\n",
    "print(len(img_files))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35cd8bb0",
   "metadata": {},
   "source": [
    "- Path模块的作用\n",
    "```shell\n",
    "> tree\n",
    ".\n",
    "└── foo\n",
    "    └── bar\n",
    "        └── test.py\n",
    "> cd foo/bar\n",
    "> cat test.py\n",
    "    # Path.resolve() #abs_path\n",
    "    from pathlib import Path\n",
    "    print(Path(__file__).parent)\n",
    "    print(Path(__file__).parent.parent)\n",
    "    \n",
    "> python test.py\n",
    "\n",
    "    . # <-- same\n",
    "    . # <-- directories\n",
    "    \n",
    "> python foo/bar/test.py\n",
    "\n",
    "    foo/bar # <-- different\n",
    "    foo     # <-- directories\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2205ea14",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.128440Z",
     "start_time": "2022-07-16T08:19:43.128429Z"
    }
   },
   "outputs": [],
   "source": [
    "# Get orientation exif tag\n",
    "for orientation in ExifTags.TAGS.keys():\n",
    "    if ExifTags.TAGS[orientation] == 'Orientation':\n",
    "        break\n",
    "\n",
    "# 此函数根据图片的信息获取图片的宽、高信息\n",
    "def exif_size(img):\n",
    "    # Returns exif-corrected PIL size\n",
    "    s = img.size  # (width, height)\n",
    "    try:\n",
    "        rotation = dict(img._getexif().items())[orientation]\n",
    "        if rotation == 6:  # rotation 270\n",
    "            s = (s[1], s[0])\n",
    "        elif rotation == 8:  # rotation 90\n",
    "            s = (s[1], s[0])\n",
    "    except:\n",
    "        pass\n",
    "    return s\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed8f048a",
   "metadata": {},
   "source": [
    "### DataLoader\n",
    "- `DataLoader`将自定义的`Dataset`根据`batch size`大小、是否`shuffle`等封装成一个`batch Size`大小的Tensor，用于后面的训练\n",
    "\n",
    "*   Dataloader本质上是一个可迭代对象，使用iter()访问，不能使用next()访问；\n",
    "*   使用iter(dataloader)返回的是一个迭代器，然后可以使用next访问；\n",
    "*   一般使用`for inputs, labels in dataloaders`进行可迭代对象的访问；"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61a26e94",
   "metadata": {},
   "source": [
    "- DataLoader参数介绍：\n",
    "```python\n",
    "DataLoader(dataset,\n",
    "           batch_size=1,\n",
    "           shuffle=False,\n",
    "           sampler=None,\n",
    "           batch_sampler=None,\n",
    "           num_workers=0,\n",
    "           collate_fn=None,   # <function default_collate>\n",
    "           pin_memory=False,\n",
    "           drop_last=False,\n",
    "           timeout=0,\n",
    "           worker_init_fn=None)\n",
    "```\n",
    "\n",
    "\n",
    "- batch_size：每个batch的大小\n",
    "- shuffle：在每个epoch开始的时候，是否对数据进行重新排序\n",
    "- num_workers：加载数据的时候使用几个子进程，0意味着所有的数据都会被load进主进程。\n",
    "- collate_fn：如何取样本，可以自己定义函数来准确地实现想要的功能\n",
    "- drop_last：告诉如何处理数据集长度除以batch_size 余下的数据。True就抛弃，否则保留\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66aed04f",
   "metadata": {},
   "source": [
    "### 示例\n",
    "```python\n",
    "# 简单使用\n",
    "dataset = MyDataset()\n",
    "dataloader = DataLoader(dataset)\n",
    "num_epoches = 100\n",
    "for epoch in range(num_epoches):\n",
    "    for img, label in dataloader:\n",
    "        ....\n",
    "```\n",
    "- ----------------------------------------------------\n",
    "```python\n",
    "\n",
    "# YOLOV3中示例\n",
    "train_dataset = YoloDataset(train_lines, input_shape, num_classes, train=True)\n",
    "val_dataset = YoloDataset(val_lines, input_shape, num_classes, train=False)\n",
    "# gen常写为train_loader\n",
    "gen = DataLoader(train_dataset,\n",
    "                 shuffle=True,\n",
    "                 batch_size=batch_size,\n",
    "                 num_workers=num_workers,\n",
    "                 pin_memory=True,\n",
    "                 drop_last=True,\n",
    "                 collate_fn=yolo_dataset_collate)\n",
    "# gen_val常写为val_loader\n",
    "gen_val = DataLoader(val_dataset,\n",
    "                     shuffle=True,\n",
    "                     batch_size=batch_size,\n",
    "                     num_workers=num_workers,\n",
    "                     pin_memory=True,\n",
    "                     drop_last=True,\n",
    "                     collate_fn=yolo_dataset_collate)\n",
    "\n",
    "for iteration, batch in enumerate(gen):\n",
    "    images, targets = batch[0], batch[1]\n",
    "    ...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a092bd6",
   "metadata": {},
   "source": [
    "### 综合使用\n",
    "\n",
    "- 数据增强手段：\n",
    "1. 裁剪(需改变bbox)\n",
    "2. 平移(需改变bbox)\n",
    "3. 改变亮度\n",
    "4. 加噪声\n",
    "5. 旋转角度(需要改变bbox)\n",
    "6. 镜像(需要改变bbox)\n",
    "7. cutout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c895f2a6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.129265Z",
     "start_time": "2022-07-16T08:19:43.129255Z"
    }
   },
   "outputs": [],
   "source": [
    "print(np.zeros((0, 5), dtype=np.float32)) # 存放label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52dc4f7e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.130133Z",
     "start_time": "2022-07-16T08:19:43.130123Z"
    },
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# %load \"/home/wssysh/data/videos/keras_yolo/yolov3_in_tf2_keras/yolodata.py\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5c17dd5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.130948Z",
     "start_time": "2022-07-16T08:19:43.130938Z"
    },
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# %load \"/home/wssysh/PycharmProjects/dataProcessing/YoloRead.py\"\n",
    "# 读取标签格式为VOC"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "958e9a2e",
   "metadata": {},
   "source": [
    "## 数据集验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb60aef8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.131851Z",
     "start_time": "2022-07-16T08:19:43.131840Z"
    }
   },
   "outputs": [],
   "source": [
    "# !pip install opencv-python\n",
    "# !pip install --upgrade pip\n",
    "# !pip install opencv-contrib-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b5d7db4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.132604Z",
     "start_time": "2022-07-16T08:19:43.132594Z"
    }
   },
   "outputs": [],
   "source": [
    "# func: 图片显示 yolo label数据与标签\n",
    "\n",
    "import numpy as np\n",
    "import cv2.cv2 as cv2\n",
    "import os\n",
    "\n",
    "# 处理好的数据集 / 初始数据\n",
    "dataset_flag = False\n",
    "\n",
    "if dataset_flag:\n",
    "    label_path = \"D://to_learn/bs/labels/train/\"\n",
    "    image_path = \"D://to_learn/bs/images/train/\"\n",
    "else:\n",
    "    label_path = \"/home/wssysh/data/datasets/temp/labels1/lb/\"\n",
    "    image_path = \"/home/wssysh/data/videos/keras_yolo/mobilenet-yolov4-lite-pytorch/VOCdevkit/VOC2007/JPEGImages\"\n",
    "#     image_path = \"/home/wssysh/data/videos/keras_yolo/bs_yolov5/dataset/VOCdevkit/VOC2007/JPEGImages\"\n",
    "\n",
    "\n",
    "# 坐标转换，原始存储的是YOLOv5格式\n",
    "# Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2]\n",
    "# where xy1=top-left, xy2=bottom-right\n",
    "def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):\n",
    "    y = np.copy(x)\n",
    "    y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw  # top left x\n",
    "    y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh  # top left y\n",
    "    y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw  # bottom right x\n",
    "    y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh  # bottom right y\n",
    "    return y\n",
    "\n",
    "\n",
    "label_lists = os.listdir(label_path)\n",
    "print(\"num of labels: \", len(label_lists))\n",
    "\n",
    "for label in label_lists:\n",
    "    # 读取labels\n",
    "    with open(os.path.join(label_path, label), 'r') as f:\n",
    "        # lines = f.readlines()\n",
    "        # print(lines)\n",
    "        lb = np.array([x.split() for x in f.read().strip().splitlines()],\n",
    "                      dtype=np.float32)  # labels\n",
    "\n",
    "    # 读取图像文件\n",
    "    img = cv2.imread(os.path.join(image_path, label.replace(\"txt\", \"jpg\")))\n",
    "    h, w = img.shape[:2]\n",
    "    lb[:, 1:] = xywhn2xyxy(lb[:, 1:], w, h, 0, 0)  # 反归一化\n",
    "\n",
    "    # 绘图\n",
    "    for _, x in enumerate(lb):\n",
    "        x_ = x.astype(int)  # np_float2int\n",
    "        class_label = \"person\" if x_[0] else \"bs\"  # class\n",
    "        cv2.rectangle(img, (x_[1], x_[2]), (x_[3], x_[4]), (0, 255, 0), 2)\n",
    "        cv2.putText(img,\n",
    "                    str(class_label), (x_[1], x_[2] - 3),\n",
    "                    fontFace=cv2.FONT_HERSHEY_SIMPLEX,\n",
    "                    fontScale=1,\n",
    "                    color=(0, 0, 255),\n",
    "                    thickness=2)\n",
    "\n",
    "    cv2.namedWindow(\"output\", cv2.WINDOW_NORMAL)\n",
    "    cv2.resizeWindow('output', 800, 600)\n",
    "    cv2.imshow(\"output\", img)\n",
    "    cv2.waitKey(70)\n",
    "\n",
    "    key = cv2.waitKey()\n",
    "    # n键 下一张\n",
    "    if key == ord(\"n\"):\n",
    "        cv2.destroyAllWindows()\n",
    "        continue\n",
    "\n",
    "    # space 输出当前图像标签\n",
    "    elif key == 32:\n",
    "        cv2.destroyAllWindows()\n",
    "        print(label.replace(\".txt\", \"\"))\n",
    "        continue\n",
    "    # Esc 退出\n",
    "    elif key == 27:\n",
    "        cv2.destroyAllWindows()\n",
    "        break\n",
    "    else:\n",
    "        cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6391b30a",
   "metadata": {},
   "source": [
    "- yolo2coco"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be1f8ec7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.133415Z",
     "start_time": "2022-07-16T08:19:43.133405Z"
    }
   },
   "outputs": [],
   "source": [
    "# r'\\cellphone_labels_cut_person_and_cellphone_total\\labels_coco_format\\annos.txt'\n",
    "\n",
    "import os\n",
    "import cv2\n",
    "\n",
    "# 原始标签路径\n",
    "originLabelsDir = \"/home/wssysh/data/datasets/temp/labels/\"\n",
    "# 转换后的文件保存路径\n",
    "saveDir = \"/home/wssysh/data/datasets/temp/labels1/\"\n",
    "# 原始标签对应的图片路径\n",
    "originImagesDir = \"/home/wssysh/data/datasets/temp/labels1/images/\"\n",
    "\n",
    "txtFileList = os.listdir(originLabelsDir)\n",
    "with open(os.path.join(saveDir, \"annos.txt\"), 'w') as fw:\n",
    "    for txtFile in txtFileList:\n",
    "        with open(os.path.join(originLabelsDir, txtFile), 'r') as fr:\n",
    "            labelList = fr.readlines()\n",
    "            for label in labelList:\n",
    "                label = label.strip().split()\n",
    "                x = float(label[1])\n",
    "                y = float(label[2])\n",
    "                w = float(label[3])\n",
    "                h = float(label[4])\n",
    "\n",
    "                # convert x,y,w,h to x1,y1,x2,y2\n",
    "                imagePath = os.path.join(originImagesDir,\n",
    "                                         txtFile.replace('txt', 'jpg'))\n",
    "                image = cv2.imread(imagePath)\n",
    "                H, W, _ = image.shape\n",
    "                x1 = (x - w / 2) * W\n",
    "                y1 = (y - h / 2) * H\n",
    "                x2 = (x + w / 2) * W\n",
    "                y2 = (y + h / 2) * H\n",
    "                # 为了与coco标签方式对，标签序号从1开始计算\n",
    "                fw.write(\n",
    "                    txtFile.replace('txt', 'jpg') + ' {} {} {} {} {}\\n'.format(\n",
    "                        int(label[0]) + 1, x1, y1, x2, y2))\n",
    "\n",
    "        print('{} done'.format(txtFile))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eed92474",
   "metadata": {},
   "source": [
    "## 数据清洗"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab261d9f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-05T13:08:45.308476Z",
     "start_time": "2023-03-05T13:08:45.279613Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# func 去除未打标图片 暂时性清洗\n",
    "# 添加标签 清洗\n",
    "import os\n",
    "import shutil\n",
    "pic_clear = True\n",
    "label_path = \"/home/wssysh/Public/\"\n",
    "img_path = \"/home/wssysh/Downloads/kaggle/archive_seeding/Keria_1904/jpg/\"\n",
    "aim_path = \"./pics/\"\n",
    "\n",
    "\n",
    "if pic_clear:\n",
    "    label_list = os.listdir(label_path)\n",
    "    # 按需替换文件后缀 和 起始名称 loop\n",
    "    img_list = []\n",
    "    for label in label_list:\n",
    "        img_list.append(label.replace(\"txt\", \"jpg\"))\n",
    "    loop = 0\n",
    "    for img in img_list:\n",
    "        loop = loop +1\n",
    "        pic_path = os.path.join(img_path, img)\n",
    "        lbs_path = (os.path.join(label_path, img)).replace(\"jpg\", \"txt\")\n",
    "        aim_path0 = os.path.join(aim_path, \"{:0>4d}.jpg\".format(loop))\n",
    "        aim_path1 = os.path.join(aim_path, \"{:0>4d}.txt\".format(loop))\n",
    "        if os.path.isfile(pic_path):\n",
    "            shutil.move(pic_path, aim_path0)\n",
    "            shutil.move(lbs_path, aim_path1)\n",
    "        else:\n",
    "            print(pic_path + \"不存在\")\n",
    "else:\n",
    "    img_list = os.listdir(img_path)\n",
    "    label_list = []\n",
    "    for img in img_list:\n",
    "        label_list.append(img.replace(\"jpg\", \"txt\"))\n",
    "    for label in label_list:\n",
    "        shutil.move(os.path.join(label_path, label), aim_path)\n",
    "\n",
    "end_num = len(os.listdir(aim_path))\n",
    "print(\"共获取有效数目：{}\".format(end_num/2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "737d0219",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:02:34.974396Z",
     "start_time": "2024-03-15T05:02:34.935457Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 将当前文件夹中的文件参考文件夹目标，复制到目标文件夹\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "# 定义源文件夹和目标文件夹的路径\n",
    "src_dir = '/home/wssysh/Data/toSucceed/videos/test/'\n",
    "d_dir = '/home/wssysh/temp/images_bk/'\n",
    "dst_dir = '/home/wssysh/temp/img/'\n",
    "\n",
    "# 遍历源文件夹中的所有文件\n",
    "for filename in os.listdir(src_dir):\n",
    "    # 检查目标文件夹中是否存在同名的文件\n",
    "    if filename in os.listdir(d_dir):\n",
    "        print(filename)\n",
    "        # 如果存在，就使用shutil模块的copy2函数将文件复制到目标文件夹\n",
    "        shutil.copy2(os.path.join(src_dir, filename), os.path.join(dst_dir, filename))\n",
    "print(\"OK!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0144d5b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.135884Z",
     "start_time": "2022-07-16T08:19:43.135872Z"
    }
   },
   "outputs": [],
   "source": [
    "# func: 根据csv文件 去除yolo dataset中不正常的图片标签及xml\n",
    "import os\n",
    "import pandas as pd\n",
    "from os.path import join as pj\n",
    "import shutil\n",
    "\n",
    "# 源地址 与 目标地址\n",
    "src_path = r\"/home/wssysh/data/datasets/bs/\"\n",
    "target_path = r\"/home/wssysh/data/datasets/temp/\"\n",
    "if not os.path.exists(target_path):\n",
    "    os.mkdir(target_path)\n",
    "\n",
    "label_path = pj(src_path, \"labels/train\")\n",
    "# image_path = pj(src_path, \"images/train\")\n",
    "image_path = \"/home/wssysh/data/videos/keras_yolo/bs_yolov5/dataset/VOCdevkit/VOC2007/JPEGImages\"\n",
    "xml_path = r\"/home/wssysh/data/videos/keras_yolo/bs_yolov5/dataset/VOCdevkit/VOC2007/Annotations/\"\n",
    "\n",
    "\n",
    "\n",
    "abnml_data = pd.read_csv(\"./abnormal.csv\", header=None)\n",
    "abnml_lists = abnml_data.iloc[:, 0]\n",
    "for abn_name in abnml_lists:\n",
    "    if not os.path.exists(pj(label_path, str(abn_name) + \".txt\")):\n",
    "        continue\n",
    "#     shutil.move(pj(label_path, str(abn_name) + \".txt\"), pj(target_path,\"labels\"))\n",
    "#     shutil.move(pj(xml_path, str(abn_name) + \".xml\"), target_path)\n",
    "#     shutil.move(pj(image_path, str(abn_name) + \".jpg\"), pj(target_path,\"labels1/images/\"))\n",
    "print(\"OK!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04442447",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-04-25T13:05:39.474075Z",
     "start_time": "2022-04-25T13:05:39.470940Z"
    }
   },
   "source": [
    "# 数据扩增"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3f3dbd4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-04T03:17:31.776422Z",
     "start_time": "2022-05-04T03:17:31.758445Z"
    }
   },
   "source": [
    "## 图像截取"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "998b07be",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-27T08:28:06.108512Z",
     "start_time": "2022-10-27T08:27:16.495451Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# func get imgs from videos\n",
    "# 建议程序运行后截取视频时选择暗色背景\n",
    "import cv2 as cv\n",
    "import os\n",
    "\n",
    "# 视频路径\n",
    "video_path = \"/home/wssysh/data/videos/bs_video/VID_20220717_193036.mp4n\"\n",
    "# 保存路径\n",
    "aim_path = \"/home/wssysh/data/videos/\"\n",
    "# 目标文件前缀 eg 210221-> 2102210001.jpg\n",
    "f_str = \"220929\"\n",
    "\n",
    "frameRate = 2  # 帧数截取间隔\n",
    "fCount = 0  # 计数器\n",
    "\n",
    "cap = cv.VideoCapture(video_path)\n",
    "\n",
    "if not os.path.exists(aim_path):\n",
    "    os.mkdir(aim_path)\n",
    "\n",
    "while (True):\n",
    "    ret, frame = cap.read()\n",
    "    fCount += 1\n",
    "    if ret:\n",
    "        if (fCount % frameRate == 0):\n",
    "            \n",
    "            # 进行反转\n",
    "            frame = cv.transpose(frame)\n",
    "            frame = cv.flip(frame,0)\n",
    "            \n",
    "            # 进行一些操作：显示截取的帧图片、保存截取帧到本地\n",
    "            cv.namedWindow(\"frame\", cv.WINDOW_NORMAL)\n",
    "            cv.resizeWindow(\"frame\", 1280, 960)\n",
    "            cv.imshow(\"frame\", frame)\n",
    "            # 延时ms，默认值为0，此时需要按键继续。\n",
    "            cv.waitKey(70)\n",
    "\n",
    "            key = cv.waitKey()\n",
    "            # n： 下一张\n",
    "            if key == ord(\"n\"):\n",
    "                cv.destroyAllWindows()\n",
    "                continue\n",
    "            # Esc： 退出\n",
    "            elif key == 27:\n",
    "                cv.destroyAllWindows()\n",
    "                break\n",
    "            # space：保存当前帧\n",
    "            elif key == 32:\n",
    "                cv.destroyAllWindows()\n",
    "                print(\"开始截取视频第：\" + str(fCount) + \" 帧\")\n",
    "                cv.imwrite(aim_path + f_str + str(\"%05d\" % fCount) + '.jpg',frame)\n",
    "                print(\"Save this frame!\")\n",
    "            # other：下一张并提示\n",
    "            else:\n",
    "                print(\"Retype please!\")\n",
    "                cv.destroyAllWindows()\n",
    "\n",
    "    else:\n",
    "        print(\"已无帧可存。\")\n",
    "        break\n",
    "cap.release()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "421c3f92",
   "metadata": {},
   "source": [
    "## 图像裁剪"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6480fe5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-15T08:33:05.333037Z",
     "start_time": "2022-11-15T08:32:59.883800Z"
    }
   },
   "outputs": [],
   "source": [
    "# func: 图片批量resize\n",
    "import os\n",
    "import cv2\n",
    "\n",
    "# 图片文件夹\n",
    "file_path = \"/home/wssysh/data/to_learn/notebooks/pics/date_20221115\"\n",
    "target_path = \"/home/wssysh/data/to_learn/notebooks/pics/date_out_20221115\"\n",
    "\n",
    "for file in os.listdir(file_path):\n",
    "    img = cv2.imread(os.path.join(file_path, file))\n",
    "    img = cv2.resize(img,(640,384))\n",
    "    \n",
    "    cv2.imwrite(os.path.join(target_path, file), img)\n",
    "    \n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11ac7ae8",
   "metadata": {},
   "source": [
    "## 其他图片转换"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e0512e0-67ba-45db-993a-48499c912ba3",
   "metadata": {},
   "source": [
    "- 通过视频帧提取\n",
    "\n",
    "```sh\n",
    "# 视频提取 质量为2\n",
    "ffmpeg -i ./VID_20240202_121313.mp4 -q:v 2 ./test/f2%04d.jpg -hide_banner\n",
    "# 间隔提取帧\n",
    "mv ./test/*7.jpg ./ori_img/\n",
    "```\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38f05c05",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.137537Z",
     "start_time": "2022-07-16T08:19:43.137527Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# module install\n",
    "# !pip install yapf\n",
    "# !pip install pillow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "673ee05c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.138539Z",
     "start_time": "2022-07-16T08:19:43.138525Z"
    }
   },
   "outputs": [],
   "source": [
    "# func：png2jpg\n",
    "from PIL import Image\n",
    "import os\n",
    "\n",
    "# png文件路径 与 目标jpg文件存储路径\n",
    "png_path = \"./capture_image/png/\"\n",
    "target_path = \"./capture_image/\"\n",
    "pngs = os.listdir(png_path)\n",
    "loop = 0\n",
    "# 循环更改\n",
    "for png in pngs:\n",
    "    img = Image.open(os.path.join(png_path, png))\n",
    "    img.save(os.path.join(target_patha, png.replace(\"png\", \"jpg\")))\n",
    "    loop += 1\n",
    "print(\"png2jpg :number：{}\".format(loop))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08cf16d8",
   "metadata": {},
   "source": [
    "### 批量调整分辨率与重命名"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33e254ac",
   "metadata": {},
   "source": [
    "- 安装ImageMagick后使用shell命令 `convert`\n",
    "\n",
    "```shell\n",
    "\n",
    "# Catch user input for file type.\n",
    "echo \"Enter the file extension for your image files:\"\n",
    "echo \"输入文件扩展名\"\n",
    "\n",
    "# Store user input in $files.\n",
    "read files\n",
    "\n",
    "loop=0\n",
    "\n",
    "# Resize and rename images.\n",
    "for img in *.$files; do\n",
    "loop=$((loop+1))\n",
    "convert -resize 512x384 \"$img\" \"resize_$loop.jpg\"\n",
    "\n",
    "done\n",
    "echo \"OK!\"\n",
    "```\n",
    "\n",
    "```shell\n",
    "# 单行模式\n",
    "for image in ./*jpg; do convert -resize 256x256 \"$image\" \"$image\" ; done\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4b534f0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-04T03:18:24.864789Z",
     "start_time": "2022-05-04T03:18:24.851803Z"
    }
   },
   "source": [
    "## 图像扩充调整"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16ce09d7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.139704Z",
     "start_time": "2022-07-16T08:19:43.139688Z"
    }
   },
   "outputs": [],
   "source": [
    "# func: 对pic_list.csv文件中的名单进行移动或删除操作\n",
    "import pandas as pd\n",
    "import os\n",
    "from PIL import Image\n",
    "from os.path import join as pj\n",
    "import shutil\n",
    "\n",
    "root_path = \"/home/wssysh/data/ceui/lbs/\"\n",
    "csvdata = pd.read_csv(\"/home/wssysh/Documents/pic_list.csv\", header=None)\n",
    "todolists = csvdata.iloc[:, 0]\n",
    "\n",
    "# 生成反转图像\n",
    "# for pic_name in todolists:\n",
    "#     img_path = pj(root_path, str(pic_name))\n",
    "#     if os.path.exists(img_path):\n",
    "#         ori_img = Image.open(img_path)\n",
    "#         r_img = ori_img.transpose(Image.FLIP_LEFT_RIGHT)\n",
    "#         r_img.save(os.path.join(\"./pic\", \"r\" + str(pic_name)))\n",
    "#     else:\n",
    "#         print(pic_name)\n",
    "#         print(\"图片路径不存在！\")\n",
    "for pic_name in todolists:\n",
    "    img_path = pj(root_path, str(pic_name) + \".txt\")\n",
    "    if os.path.exists(img_path):\n",
    "        shutil.move(img_path, pj(\"./lbs/\", str(pic_name)))\n",
    "    else:\n",
    "        print(pic_name)\n",
    "        print(\"图片路径不存在！\")\n",
    "\n",
    "print(\"OK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "277bc4e6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:18:56.182166Z",
     "start_time": "2024-03-15T05:18:49.216151Z"
    },
    "code_folding": [],
    "tags": []
   },
   "outputs": [],
   "source": [
    "# func：进行图片左右翻转 及 所有图片的颜色调整\n",
    "# -*- coding: UTF-8 -*-\n",
    "\n",
    "from PIL import Image\n",
    "from PIL import ImageEnhance\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "# 文件路径 与 目标 jpg文件存储路径\n",
    "# label_path = \"/home/wssysh/Public/temp/imgs/\"\n",
    "img_path = \"/home/wssysh/temp/lanedata/imgs/\"\n",
    "aim_path = \"/home/wssysh/temp/lanedata/imgsc/\"\n",
    "# aim_path = \"/home/wssysh/temp/imgcs/\"\n",
    "enh_target_path = aim_path\n",
    "\n",
    "# 翻转后存储目录\n",
    "# LR_target_path = os.path.join(aim_path, \"LR\")\n",
    "# TB_target_path = os.path.join(aim_path, \"TB\")\n",
    "# rotate ccw 顺逆时针旋转 仅适用于 宽高相等\n",
    "# ccw_target_path = os.path.join(aim_path, \"ccw\")\n",
    "\n",
    "# 文件夹不存在时新建\n",
    "# if not os.path.exists(\"../output/\"):\n",
    "#     os.mkdir(\"../output/\")\n",
    "if not os.path.exists(enh_target_path):\n",
    "    os.mkdir(enh_target_path)\n",
    "# if not os.path.exists(LR_target_path):\n",
    "#     os.mkdir(LR_target_path)\n",
    "# if not os.path.exists(TB_target_path):\n",
    "#     os.mkdir(TB_target_path)\n",
    "# if not os.path.exists(ccw_target_path):\n",
    "#     os.mkdir(ccw_target_path)\n",
    "\n",
    "\n",
    "# 调整图像饱和度、亮度、对比度、锐度\n",
    "def enh_color(img):\n",
    "    # 调整图像的饱和度\n",
    "    random_factor = (\n",
    "        np.random.randint(6, 8) / 10.0\n",
    "        if np.random.randint(1, 3) == 1\n",
    "        else np.random.randint(12, 15) / 10.0\n",
    "    )\n",
    "    color_image = ImageEnhance.Color(img).enhance(random_factor)\n",
    "\n",
    "    # 调整图像的亮度\n",
    "    random_factor = (\n",
    "        np.random.randint(6, 8) / 10.0\n",
    "        if np.random.randint(1, 3) == 1\n",
    "        else np.random.randint(11, 13) / 10.0\n",
    "    )\n",
    "    brightness_image = ImageEnhance.Brightness(color_image).enhance(random_factor)\n",
    "    # 调整图像对比度\n",
    "    random_factor = (\n",
    "        np.random.randint(6, 8) / 10.0\n",
    "        if np.random.randint(1, 3) == 1\n",
    "        else np.random.randint(12, 16) / 10.0\n",
    "    )\n",
    "    contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor)\n",
    "\n",
    "    # 调整图像锐度\n",
    "    random_factor = (\n",
    "        np.random.randint(15, 31) / 10.0\n",
    "        if np.random.randint(1, 3) == 1\n",
    "        else np.random.randint(6, 8) / 10.0\n",
    "    )  # 随机因子\n",
    "    result = ImageEnhance.Sharpness(contrast_image).enhance(random_factor)\n",
    "    return result\n",
    "\n",
    "\n",
    "# 读取文件并计数操作原始图像数目\n",
    "loop = 0\n",
    "imgs = os.listdir(img_path)\n",
    "\n",
    "# 保存\n",
    "for img in imgs:\n",
    "    ori_img = Image.open(os.path.join(img_path, img))\n",
    "    # 颜色调整\n",
    "    enh_img = enh_color(ori_img)\n",
    "    enh_img.save(os.path.join(enh_target_path, \"CC\" + img))\n",
    "\n",
    "    # 逆时针旋转90 ccw\n",
    "    # ccw_img = ori_img.rotate(angle=90, expand=True)\n",
    "    # ccw_img.save(os.path.join(ccw_target_path, \"ccw\" + img))\n",
    "\n",
    "    # 上下 TB\n",
    "    # TB_img = ori_img.transpose(Image.FLIP_TOP_BOTTOM)\n",
    "    # TB_img.save(os.path.join(TB_target_path, \"TB\" + img))\n",
    "\n",
    "    # 左右翻转 LR\n",
    "    # LR_img = ori_img.transpose(Image.FLIP_LEFT_RIGHT)\n",
    "    # LR_img.save(os.path.join(aim_path, \"LR\" + img))\n",
    "\n",
    "    # # 批量重命名\n",
    "    loop += 1\n",
    "\n",
    "print(\"img number：{}\".format(loop))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b3903d0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.142427Z",
     "start_time": "2022-07-16T08:19:43.142411Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# %load \"/home/wssysh/data/videos/keras_yolo/bs_yolov5/dataset/VOCdevkit/img_flip.py\"\n",
    "# func：进行图片左右翻转 及 所有图片的颜色调整\n",
    "from PIL import Image\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "# 文件路径 与 目标jpg文件存储路径\n",
    "img_path = \"/home/wssysh/data/labels/pics/\"\n",
    "\n",
    "flip_target_path = \"/home/wssysh/data/labels/pics_tb/\"\n",
    "# flip_target_path = \"/home/wssysh/data/labels/pics_lr/\"\n",
    "\n",
    "if not os.path.exists(flip_target_path):\n",
    "    os.mkdir(flip_target_path)\n",
    "\n",
    "\n",
    "# 读取文件并计数操作原始图像数目\n",
    "loop = 0\n",
    "imgs = os.listdir(img_path)\n",
    "\n",
    "# 保存\n",
    "for img in imgs:\n",
    "    ori_img = Image.open(os.path.join(img_path, img))\n",
    "    # 左右翻转\n",
    "    r_img = ori_img.transpose(Image.FLIP_TOP_BOTTOM)\n",
    "    r_img.save(os.path.join(flip_target_path, \"tb\" + img))\n",
    "    \n",
    "    loop += 1\n",
    "\n",
    "print(\"img number：{}\".format(loop))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6c38a58",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:27:10.553515Z",
     "start_time": "2024-03-15T05:27:09.773748Z"
    }
   },
   "outputs": [],
   "source": [
    "# 关键点检测对标签进行处理\n",
    "import os\n",
    "import pandas as pd\n",
    "from os.path import join as pj\n",
    "# 标签地址 与 目标地址\n",
    "label_path = \"/home/wssysh/temp/lanedata/labels/\"\n",
    "out_label_path = \"/home/wssysh/temp/lanedata/tempyolo/\"\n",
    "\n",
    "if not os.path.exists(out_label_path):\n",
    "    os.mkdir(out_label_path)\n",
    "\n",
    "labels = os.listdir(label_path)\n",
    "for label in labels:\n",
    "    # debug:\n",
    "    print(label+\" 正在处理\")\n",
    "    \n",
    "    #sep：每个数直接是用空格隔开的\n",
    "    label_data = pd.read_csv(pj(label_path, label), header=None, sep=\" \")    \n",
    "    \n",
    "    #逆时针旋转90\n",
    "    # op_str = 'ccw'\n",
    "    # x=label_data.iloc[:,1].copy()\n",
    "    # y=label_data.iloc[:,3].copy()\n",
    "    # label_data.iloc[:, 1] = label_data.iloc[:, 2]\n",
    "    # label_data.iloc[:, 2] = -(x - 1.0)\n",
    "    # label_data.iloc[:, 3] = label_data.iloc[:, 4]\n",
    "    # label_data.iloc[:, 4] = y\n",
    "    \n",
    "    # 反转的标签数据处理 注意切片广播机制 1左右翻转 2上下翻转\n",
    "    op_str = 'LR'\n",
    "    i_num=0\n",
    "    label_data.iloc[:, 1] = -(label_data.iloc[:, 1] - 1.0)\n",
    "\n",
    "    # 处理关键点\n",
    "    for ind_num in range(5,50,3):\n",
    "        i_num=i_num+1.0\n",
    "        # print(ind_num,i_num)\n",
    "        # label_data.iloc[:, 5] = -(label_data.iloc[:, 5] - 1.0)\n",
    "        # 对单框适用\n",
    "        # print(label_data.iloc[0, ind_num])\n",
    "        if label_data.iloc[0, ind_num]!=0:\n",
    "            print(i_num)\n",
    "            label_data.iloc[0, ind_num] = -(label_data.iloc[0, ind_num] - 1.0)\n",
    "        # 实现水平翻转后的坐标调整\n",
    "    # break\n",
    "    \n",
    "    # 保存"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f8465e6",
   "metadata": {},
   "source": [
    "## 预打标\n",
    "\n",
    "- 通过预训练的模型扩增数据集的标签\n",
    "- 若需要voc格式的标签文件则用后面的yolo2voc进行转换\n",
    "- 转换前在yolo标签文件夹指定好 classes.txt, 详情请自行查阅资料"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90a28f92",
   "metadata": {},
   "outputs": [],
   "source": [
    "# func: yolo2yolo_txt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from yolo import YOLO\n",
    "import os\n",
    "from os.path import join as pj\n",
    "\n",
    "# 在YOLO detect_image 中添加flag\n",
    "# if box_f:\n",
    "#     return top_label, top_boxes\n",
    "\n",
    "\n",
    "# 坐标转换\n",
    "# 若出现位置不对称问题需自行分析数据含义并修改\n",
    "def xyxy2xywh(size, box):\n",
    "    dw = 1. / size[1]\n",
    "    dh = 1. / size[0]\n",
    "    x = (box[0] + box[2]) / 2.0\n",
    "    y = (box[1] + box[3]) / 2.0\n",
    "    w = box[2] - box[0]\n",
    "    h = box[3] - box[1]\n",
    "    x = x * dh\n",
    "    w = w * dh\n",
    "    y = y * dw\n",
    "    h = h * dw\n",
    "    return (y, x, h, w)\n",
    "\n",
    "\n",
    "# 实例化YOLO对象\n",
    "yolo = YOLO()\n",
    "# 图片路径 只能包含jpg文件\n",
    "ImagesDir = \"VOCdevkit/VOC2007/JPEGImages\"\n",
    "# 保存路径\n",
    "saveDir = \"model_data/new_yolo_label\"\n",
    "\n",
    "# img_path = \"VOCdevkit/VOC2007/JPEGImages\"\n",
    "# print([x for x in xs])\n",
    "\n",
    "\n",
    "# 单个文件转换：图片，输出路径\n",
    "def pre_det2yolo(img_path, img_name, out_dir):\n",
    "    image = Image.open(pj(img_path, img_name))\n",
    "    image_shape = np.array(np.shape(image)[0:2])\n",
    "    label, box_ori = yolo.detect_image(image, box_f=True)\n",
    "    yolo_txt = open(pj(out_dir, img_name.replace(\"jpg\", \"txt\")), 'w')\n",
    "    # print(image_shape, box_ori)\n",
    "    loop = 0\n",
    "    for lb in label:\n",
    "        yolo_box = xyxy2xywh(image_shape, box_ori[loop])\n",
    "        yolo_txt.write(\n",
    "            str(lb) + \n",
    "            \" \" + str(yolo_box[0]) + \n",
    "            \" \" + str(yolo_box[1]) + \n",
    "            \" \" + str(yolo_box[2]) + \n",
    "            \" \" + str(yolo_box[3]) + '\\n')\n",
    "        loop += 1\n",
    "        # print(lb, yolo_box)\n",
    "\n",
    "    yolo_txt.close()\n",
    "\n",
    "\n",
    "# 批量处理文件\n",
    "img_list = os.listdir(ImagesDir)\n",
    "loop_1 = 0\n",
    "print(\"预测生成中,请稍后...\")\n",
    "for img in img_list:\n",
    "    pre_det2yolo(ImagesDir, img, saveDir)\n",
    "    loop_1 = loop_1 + 1\n",
    "    # print(loop_1)\n",
    "\n",
    "print(\"共生成标签文件:%d\" % loop_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "469de8be",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T13:55:49.569733Z",
     "start_time": "2024-03-15T13:55:49.355069Z"
    }
   },
   "outputs": [],
   "source": [
    "# YOLO POSE to label\n",
    "# 本脚本适用于单框关键点检测\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "from ultralytics import YOLO\n",
    "from PIL import Image, ImageDraw\n",
    "\n",
    "def pre_pose2yolo(img_path, out_dir):\n",
    "    model = YOLO(\"/home/wssysh/Data/toSucceed/Lanedet/best.pt\")\n",
    "    imgs = os.listdir(img_path)\n",
    "    for img in imgs:\n",
    "        results = model.predict(source=os.path.join(img_path, img))\n",
    "        for r in results:\n",
    "            label_list = []\n",
    "            # this line is changed\n",
    "            # print(r.boxes)\n",
    "            boxes_val = r.boxes.xywhn.cpu().numpy()  # get the keypoints\n",
    "            # print(boxes_val.flatten())\n",
    "            label_list.append(0)\n",
    "            for val in boxes_val.flatten():\n",
    "                label_list.append(val)\n",
    "            # print(label_list)\n",
    "            # print(r.keypoints)\n",
    "            keypoints_val = r.keypoints.xyn.cpu().numpy()  # get the keypoints\n",
    "            for xy in keypoints_val[0]:\n",
    "                if xy[0] != 0 and xy[1] != 0:\n",
    "                    for val in xy:\n",
    "                        label_list.append(val)\n",
    "                    label_list.append(2)\n",
    "                else:\n",
    "                    for val in xy:\n",
    "                        label_list.append(val)\n",
    "                    label_list.append(0)\n",
    "        \n",
    "        # print(label_list)\n",
    "        # 写出文件\n",
    "        with open(os.path.join(out_dir,os.path.splitext(img)[0] + \".txt\"), \"w\") as f:\n",
    "            f.write(\" \".join(map(str, label_list)))\n",
    "    print(\"OK!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa29f95b-2b4f-40d3-bd37-634e617a5ee6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "img_path = \"/home/wssysh/temp/lanedata/imgs/\"\n",
    "out_dir = \"/home/wssysh/temp/lanedata/lbs/\"\n",
    "pre_pose2yolo(img_path, out_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4030ad5d",
   "metadata": {},
   "source": [
    "## 标签处理"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35cf9f86",
   "metadata": {},
   "source": [
    "### 图像扩增对应标签处理\n",
    "\n",
    "- 只适用yolo的txt格式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "407a9136",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load ~/test.py\n",
    "import os\n",
    "import pandas as pd\n",
    "from os.path import join as pj\n",
    "# 标签地址 与 目标地址\n",
    "label_path = \"/home/wssysh/temp/lanedata/labels/\"\n",
    "out_label_path = \"/home/wssysh/temp/lanedata/tempyolo/\"\n",
    "\n",
    "if not os.path.exists(out_label_path):\n",
    "    os.mkdir(out_label_path)\n",
    "\n",
    "labels = os.listdir(label_path)\n",
    "for label in labels:\n",
    "    # debug:\n",
    "    print(label + \" 正在处理\")\n",
    "\n",
    "    # sep：每个数直接是用空格隔开的\n",
    "    label_data = pd.read_csv(pj(label_path, label), header=None, sep=\" \")\n",
    "\n",
    "    # 逆时针旋转90\n",
    "    # op_str = 'ccw'\n",
    "    # x=label_data.iloc[:,1].copy()\n",
    "    # y=label_data.iloc[:,3].copy()\n",
    "    # label_data.iloc[:, 1] = label_data.iloc[:, 2]\n",
    "    # label_data.iloc[:, 2] = -(x - 1.0)\n",
    "    # label_data.iloc[:, 3] = label_data.iloc[:, 4]\n",
    "    # label_data.iloc[:, 4] = y\n",
    "\n",
    "    # 反转的标签数据处理 注意切片广播机制 1左右翻转 2上下翻转\n",
    "    op_str = \"LR\"\n",
    "\n",
    "    # 处理锚框  ---------------------------------\n",
    "    label_data.iloc[:, 1] = -(label_data.iloc[:, 1] - 1.0)\n",
    "\n",
    "    # 处理关键点 ---------------------------------\n",
    "    # 此处对单框适用 ，若多框请进行循环处理\n",
    "    for ind_num in range(5, 50, 3):\n",
    "        # label_data.iloc[:, 5] = -(label_data.iloc[:, 5] - 1.0)\n",
    "        # 实现水平翻转后的坐标调整\n",
    "        if label_data.iloc[0, ind_num+2] != 0:\n",
    "            label_data.iloc[0, ind_num] = -(label_data.iloc[0, ind_num] - 1.0)\n",
    "    # 实现类别翻转----- 特殊标记\n",
    "    temp_num = 0\n",
    "    for ind_num in range(5, 20, 3):\n",
    "        for i in range(0, 3, 1):\n",
    "            temp_num = label_data.iloc[0, ind_num + i]\n",
    "            label_data.iloc[0, ind_num + i] = label_data.iloc[0, ind_num + i + 30]\n",
    "            label_data.iloc[0, ind_num + i + 30] = temp_num\n",
    "\n",
    "    # 保存\n",
    "    label_data.to_csv(\n",
    "        pj(out_label_path, op_str + label),\n",
    "        sep=\" \",\n",
    "        float_format=\"%.5f\",\n",
    "        index=None,\n",
    "        header=None,\n",
    "    )\n",
    "    # debug:\n",
    "    # break\n",
    "print(\"OK!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c66812fd",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.144612Z",
     "start_time": "2022-07-16T08:19:43.144600Z"
    }
   },
   "outputs": [],
   "source": [
    "i_num = 0\n",
    "for ind_num in range(5, 50, 3):\n",
    "    i_num = i_num + 1\n",
    "    print(i_num, ind_num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa618c00",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.145423Z",
     "start_time": "2022-07-16T08:19:43.145412Z"
    }
   },
   "outputs": [],
   "source": [
    "from PIL import Image,ImageEnhance\n",
    "ori_img = Image.open(\"/home/wssysh/22042406320.jpg\")\n",
    "ori_img.show()\n",
    "color_image = ImageEnhance.Color(ori_img).enhance(1.3)\n",
    "bright_image = ImageEnhance.Brightness(color_image).enhance(1.3)\n",
    "contrast_image = ImageEnhance.Contrast(bright_image).enhance(0.7)\n",
    "contrast_image.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "949706dd",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.146326Z",
     "start_time": "2022-07-16T08:19:43.146315Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.randint(1,3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "572d16bd",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.147317Z",
     "start_time": "2022-07-16T08:19:43.147296Z"
    }
   },
   "outputs": [],
   "source": [
    "random_factor = np.random.randint(3, 7) / 10. if np.random.randint(1,3) == 1 else np.random.randint(18,22) / 10.\n",
    "print(random_factor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56203f72",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-05T13:25:03.787122Z",
     "start_time": "2023-03-05T13:25:03.331940Z"
    }
   },
   "outputs": [],
   "source": [
    "# 对水平flip图片进行标签处理 上下的需iloc[:, 2]\n",
    "# 使用labelimg进行检察,防止出错\n",
    "# 生成数据标签前缀加r eg: r22040403435.txt\n",
    "\n",
    "import os\n",
    "import pandas as pd\n",
    "from os.path import join as pj\n",
    "\n",
    "# 标签地址 与 目标地址\n",
    "label_path = r\"./lbs/\"\n",
    "r_label_path = label_path.replace(\"lbs\", \"rlabels\")\n",
    "\n",
    "if not os.path.exists(r_label_path):\n",
    "    os.mkdir(r_label_path)\n",
    "\n",
    "labels = os.listdir(label_path)\n",
    "for label in labels:\n",
    "    # print(label + \" 正在处理...\")\n",
    "\n",
    "    # sep：每个数直接是用空格隔开的\n",
    "    label_data = pd.read_csv(pj(label_path, label), header=None, sep=\" \")\n",
    "\n",
    "    # 反转的标签数据处理 注意切片广播机制\n",
    "    label_data.iloc[:, 1] = -(label_data.iloc[:, 1] - 1.0)\n",
    "    # 保存\n",
    "    label_data.to_csv(pj(r_label_path, \"LR\" + label), sep=\" \", index=None, header=None)\n",
    "print(\"OK!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bf0e601",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.150033Z",
     "start_time": "2022-07-16T08:19:43.150018Z"
    }
   },
   "outputs": [],
   "source": [
    "# 标签批量复制并重命名\n",
    "import os\n",
    "from os.path import join as pj\n",
    "import shutil\n",
    "\n",
    "# 标签地址 与 目标地址\n",
    "# label_path = r\"/home/wssysh/data/datasets/temp/labels1/lb/\"\n",
    "label_path = \"/home/wssysh/data/to_learn/notebooks/output/r_labels/\"\n",
    "out_path = r\"../output/\"\n",
    "\n",
    "# 此处 修改前缀\n",
    "prefix_str = \"enh_\"\n",
    "enh_label_path = pj(out_path, prefix_str + \"r_labels\")\n",
    "\n",
    "if not os.path.exists(enh_label_path):\n",
    "    os.mkdir(enh_label_path)\n",
    "\n",
    "file_names = os.listdir(label_path)\n",
    "for name in file_names: \n",
    "    shutil.copy(pj(label_path, name), pj(enh_label_path, prefix_str + name))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3eebd878",
   "metadata": {},
   "source": [
    "### 标签格式转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc260c02",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-18T08:56:10.272716Z",
     "start_time": "2023-12-18T08:56:09.217256Z"
    },
    "code_folding": [],
    "tags": []
   },
   "outputs": [],
   "source": [
    "# func: voc2yolo\n",
    "# xml2txt\n",
    "#!/usr/bin/env python3\n",
    "# -*- coding: utf-8 -*-\n",
    "\n",
    "import xml.etree.ElementTree as ET\n",
    "import os\n",
    "\n",
    "# 输入路径与输出路径\n",
    "xml_path = '/home/wssysh/Data/packs/src_git/yolov4-tiny-pytorch/VOCdevkit/VOC2007/Annotations/'\n",
    "outpath = \"/home/wssysh/Public/temp/\"\n",
    "\n",
    "# 注意此处的类别与xml中对应\n",
    "classes = [\"seeding\"]  # 类别\n",
    "\n",
    "\n",
    "def convert(size, box):\n",
    "    dw = 1. / size[0]\n",
    "    dh = 1. / size[1]\n",
    "    x = (box[0] + box[1]) / 2.0\n",
    "    y = (box[2] + box[3]) / 2.0\n",
    "    w = box[1] - box[0]\n",
    "    h = box[3] - box[2]\n",
    "    x = x * dw\n",
    "    w = w * dw\n",
    "    y = y * dh\n",
    "    h = h * dh\n",
    "    return (x, y, w, h)\n",
    "\n",
    "\n",
    "def convert_annotation(label_name):\n",
    "    xml_file = open(xml_path + label_name, encoding='UTF-8')\n",
    "    txt_file = open(outpath + label_name.replace(\"xml\", \"txt\"), 'w')\n",
    "    tree = ET.parse(xml_file)\n",
    "    root = tree.getroot()\n",
    "    size = root.find('size')\n",
    "    w = int(size.find('width').text)\n",
    "    h = int(size.find('height').text)\n",
    "    # print(w,h)\n",
    "\n",
    "    for obj in root.iter('object'):\n",
    "        cls = obj.find('name').text\n",
    "        # print(cls)\n",
    "        if cls not in classes:\n",
    "            continue\n",
    "        cls_id = classes.index(cls)\n",
    "        xmlbox = obj.find('bndbox')\n",
    "        ori_b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),\n",
    "             float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))\n",
    "        bb = convert((w, h), ori_b)\n",
    "        txt_file.write(str(cls_id) + \" \" + \" \".join([str(a) for a in bb]) + '\\n')\n",
    "\n",
    "\n",
    "# xml list\n",
    "img_xmls = os.listdir(xml_path)\n",
    "for img_xml in img_xmls:\n",
    "    convert_annotation(img_xml)\n",
    "print(\"OK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "763fc01b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-19T02:37:15.590971Z",
     "start_time": "2023-12-19T02:37:08.276148Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# func: yolo2voc\n",
    "# 注意切换路径，nonetype一般是读取失败\n",
    "# AttributeError: 'NoneType' object has no attribute 'shape'\n",
    "\n",
    "# txtPath, xmlPath, picPath txt所在文件夹路径，xml文件保存路径，图片所在文件夹路径\n",
    "\n",
    "from xml.dom.minidom import Document\n",
    "import os\n",
    "import cv2\n",
    " \n",
    " \n",
    "def makexml(picPath, txtPath, xmlPath):  # txt所在文件夹路径，xml文件保存路径，图片所在文件夹路径\n",
    "    \"\"\"此函数用于将yolo格式txt标注文件转换为voc格式xml标注文件\n",
    "    在自己的标注图片文件夹下建三个子文件夹，分别命名为imgs,lbs(yolo labels),data(to save xml files)\n",
    "    \"\"\"\n",
    "    dic = {'0': \"seedling\",  # 创建字典用来对类型进行转换\n",
    "           '1': \"weed\",  # 此处的字典要与自己的classes.txt文件中的类对应，且顺序要一致\n",
    "           }\n",
    "    files = os.listdir(txtPath)\n",
    "    for i, name in enumerate(files):\n",
    "        # Debug!!\n",
    "        # print(name)\n",
    "        xmlBuilder = Document()\n",
    "        annotation = xmlBuilder.createElement(\"annotation\")  # 创建annotation标签\n",
    "        xmlBuilder.appendChild(annotation)\n",
    "        txtFile = open(txtPath + name)\n",
    "        txtList = txtFile.readlines()\n",
    "        img = cv2.imread(picPath + name[0:-4] + \".jpg\")\n",
    "        Pheight, Pwidth, Pdepth = img.shape\n",
    " \n",
    "        folder = xmlBuilder.createElement(\"folder\")  # folder标签\n",
    "        foldercontent = xmlBuilder.createTextNode(\"driving_annotation_dataset\")\n",
    "        folder.appendChild(foldercontent)\n",
    "        annotation.appendChild(folder)  # folder标签结束\n",
    " \n",
    "        filename = xmlBuilder.createElement(\"filename\")  # filename标签\n",
    "        filenamecontent = xmlBuilder.createTextNode(name[0:-4] + \".jpg\")\n",
    "        filename.appendChild(filenamecontent)\n",
    "        annotation.appendChild(filename)  # filename标签结束\n",
    " \n",
    "        size = xmlBuilder.createElement(\"size\")  # size标签\n",
    "        width = xmlBuilder.createElement(\"width\")  # size子标签width\n",
    "        widthcontent = xmlBuilder.createTextNode(str(Pwidth))\n",
    "        width.appendChild(widthcontent)\n",
    "        size.appendChild(width)  # size子标签width结束\n",
    " \n",
    "        height = xmlBuilder.createElement(\"height\")  # size子标签height\n",
    "        heightcontent = xmlBuilder.createTextNode(str(Pheight))\n",
    "        height.appendChild(heightcontent)\n",
    "        size.appendChild(height)  # size子标签height结束\n",
    " \n",
    "        depth = xmlBuilder.createElement(\"depth\")  # size子标签depth\n",
    "        depthcontent = xmlBuilder.createTextNode(str(Pdepth))\n",
    "        depth.appendChild(depthcontent)\n",
    "        size.appendChild(depth)  # size子标签depth结束\n",
    " \n",
    "        annotation.appendChild(size)  # size标签结束\n",
    " \n",
    "        for j in txtList:\n",
    "            oneline = j.strip().split(\" \")\n",
    "            object = xmlBuilder.createElement(\"object\")  # object 标签\n",
    "            picname = xmlBuilder.createElement(\"name\")  # name标签\n",
    "            namecontent = xmlBuilder.createTextNode(dic[oneline[0]])\n",
    "            picname.appendChild(namecontent)\n",
    "            object.appendChild(picname)  # name标签结束\n",
    " \n",
    "            pose = xmlBuilder.createElement(\"pose\")  # pose标签\n",
    "            posecontent = xmlBuilder.createTextNode(\"Unspecified\")\n",
    "            pose.appendChild(posecontent)\n",
    "            object.appendChild(pose)  # pose标签结束\n",
    " \n",
    "            truncated = xmlBuilder.createElement(\"truncated\")  # truncated标签\n",
    "            truncatedContent = xmlBuilder.createTextNode(\"0\")\n",
    "            truncated.appendChild(truncatedContent)\n",
    "            object.appendChild(truncated)  # truncated标签结束\n",
    " \n",
    "            difficult = xmlBuilder.createElement(\"difficult\")  # difficult标签\n",
    "            difficultcontent = xmlBuilder.createTextNode(\"0\")\n",
    "            difficult.appendChild(difficultcontent)\n",
    "            object.appendChild(difficult)  # difficult标签结束\n",
    " \n",
    "            bndbox = xmlBuilder.createElement(\"bndbox\")  # bndbox标签\n",
    "            xmin = xmlBuilder.createElement(\"xmin\")  # xmin标签\n",
    "            mathData = int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth)\n",
    "            xminContent = xmlBuilder.createTextNode(str(mathData))\n",
    "            xmin.appendChild(xminContent)\n",
    "            bndbox.appendChild(xmin)  # xmin标签结束\n",
    " \n",
    "            ymin = xmlBuilder.createElement(\"ymin\")  # ymin标签\n",
    "            mathData = int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight)\n",
    "            yminContent = xmlBuilder.createTextNode(str(mathData))\n",
    "            ymin.appendChild(yminContent)\n",
    "            bndbox.appendChild(ymin)  # ymin标签结束\n",
    " \n",
    "            xmax = xmlBuilder.createElement(\"xmax\")  # xmax标签\n",
    "            mathData = int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth)\n",
    "            xmaxContent = xmlBuilder.createTextNode(str(mathData))\n",
    "            xmax.appendChild(xmaxContent)\n",
    "            bndbox.appendChild(xmax)  # xmax标签结束\n",
    " \n",
    "            ymax = xmlBuilder.createElement(\"ymax\")  # ymax标签\n",
    "            mathData = int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight)\n",
    "            ymaxContent = xmlBuilder.createTextNode(str(mathData))\n",
    "            ymax.appendChild(ymaxContent)\n",
    "            bndbox.appendChild(ymax)  # ymax标签结束\n",
    " \n",
    "            object.appendChild(bndbox)  # bndbox标签结束\n",
    " \n",
    "            annotation.appendChild(object)  # object标签结束\n",
    " \n",
    "        f = open(xmlPath + name[0:-4] + \".xml\", 'w')\n",
    "        xmlBuilder.writexml(f, indent='\\t', newl='\\n', addindent='\\t', encoding='UTF-8')\n",
    "        f.close()\n",
    " \n",
    "if __name__ == \"__main__\":\n",
    "    # os.system(\"mkdir -vp ./dataset/VOCdevkit/VOC2007/Annotations/\")\n",
    "    # os.system(\"mkdir ./dataset/VOCdevkit/VOC2007/JPEGImages/\")\n",
    "    # os.system(\"cp /home/wssysh/data/datasets/bs/images/valid/* ./dataset/VOCdevkit/VOC2007/JPEGImages/\")\n",
    "    # 注意图片后缀，对上面的代码进行修改，如jpg jpeg png\n",
    "    rootPath = \"/home/wssysh/Public/temp/\" # 根目录\n",
    "    picPath = rootPath + \"pics/\"  # 图片所在文件夹路径，后面的/一定要带上\n",
    "    txtPath = rootPath + \"lbs/\"  # 后面的/一定要带上\n",
    "    \n",
    "    xmlPath = rootPath + \"xmls/\"  # 后面的/一定要带上\n",
    "    if not os.path.exists(xmlPath):\n",
    "        os.mkdir(xmlPath)\n",
    "    \n",
    "    makexml(picPath, txtPath, xmlPath)\n",
    "    print(\"OK!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1807c511",
   "metadata": {},
   "source": [
    "#### YOLO txt to labelme json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a88491b1-0471-4dd7-80c1-4e29ca276c81",
   "metadata": {},
   "source": [
    "<img src=\"./saved_pic/yololbs.jpg\" style=\"width:400px\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "274200cb-ca57-42fc-89f7-ae7012074e74",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"./saved_pic/yololbs.jpg\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53af3e33",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:04:46.845227Z",
     "start_time": "2024-03-15T10:04:46.836919Z"
    }
   },
   "outputs": [],
   "source": [
    "# txt2json\n",
    "# 注意本代码只处理了单框多关键点的问题\n",
    "# 类别需自行修改\n",
    "\n",
    "import os\n",
    "import json\n",
    "\n",
    "# import base64\n",
    "# import cv2\n",
    "# from ultralytics.utils.ops import crop_mask, xywh2xyxy, xyxy2xywh\n",
    "\n",
    "\n",
    "def read_txt_file(txt_file):\n",
    "    with open(txt_file, 'r') as f:\n",
    "        lines = f.readlines()\n",
    "    data = []\n",
    "    for line in lines:\n",
    "        line = line.strip().split()\n",
    "        class_name = line[0]\n",
    "        bbox = [coord for coord in line[1:]]\n",
    "        data.append({'class_name': class_name, 'bbox': bbox})\n",
    "    return data\n",
    "\n",
    "\n",
    "def convert_to_labelme(data, image_name, image_size):\n",
    "    labelme_data = {\n",
    "        'version': '5.4.1',\n",
    "        'flags': {},\n",
    "        'shapes': [],\n",
    "        'imagePath': image_name,\n",
    "        'imageData': None,\n",
    "        'imageHeight': image_size[1],\n",
    "        'imageWidth': image_size[0]\n",
    "    }\n",
    "\n",
    "    for obj in data:\n",
    "        dx = obj['bbox'][0]\n",
    "        dy = obj['bbox'][1]\n",
    "        dw = obj['bbox'][2]\n",
    "        dh = obj['bbox'][3]\n",
    "\n",
    "        w = eval(dw) * image_size[0]\n",
    "        h = eval(dh) * image_size[1]\n",
    "        center_x = eval(dx) * image_size[0]\n",
    "        center_y = eval(dy) * image_size[1]\n",
    "        x1 = center_x - w / 2\n",
    "        y1 = center_y - h / 2\n",
    "        x2 = center_x + w / 2\n",
    "        y2 = center_y + h / 2\n",
    "\n",
    "        # rect\n",
    "        if obj['class_name'] == '0':  #判断对应的标签名称，写入json文件中\n",
    "            label = str('roi')\n",
    "        else:\n",
    "            label = obj['class_name']\n",
    "        # 锚框类别自动调整\n",
    "\n",
    "        shape_data = {\n",
    "            'label': label,\n",
    "            'points': [[x1, y1], [x2, y2]],\n",
    "            'group_id': None,\n",
    "            'shape_type': 'rectangle',\n",
    "            'flags': {}\n",
    "        }\n",
    "        labelme_data['shapes'].append(shape_data)\n",
    "        # points 请修改：共 15 个点\n",
    "        # 注意关键点类别\n",
    "        labels = [\n",
    "            'llane0', 'llane1', 'llane2', 'llane3', 'llane4', 'lane0', 'lane1',\n",
    "            'lane2', 'lane3', 'lane4', 'rlane0', 'rlane1', 'rlane2', 'rlane3',\n",
    "            'rlane4'\n",
    "        ]\n",
    "        idx = 0\n",
    "        for ind_num in range(4, 47, 3):\n",
    "            if obj['bbox'][ind_num + 2] != \"0\":\n",
    "                pointx = eval(obj['bbox'][ind_num]) * image_size[0]\n",
    "                pointy = eval(obj['bbox'][ind_num + 1]) * image_size[1]\n",
    "                shape_data = {\n",
    "                    'label': labels[idx],\n",
    "                    'points': [[pointx, pointy]],\n",
    "                    'group_id': None,\n",
    "                    'shape_type': 'point',\n",
    "                    'flags': {}\n",
    "                }\n",
    "                labelme_data['shapes'].append(shape_data)\n",
    "            idx = idx + 1\n",
    "    return labelme_data\n",
    "\n",
    "\n",
    "def save_labelme_json(labelme_data, image_path, output_file):\n",
    "    #     with open(image_path, 'rb') as f:\n",
    "    #         image_data = f.read()\n",
    "    labelme_data['imageData'] = None  #base64.b64encode(image_data).decode('utf-8')\n",
    "\n",
    "    with open(output_file, 'w') as f:\n",
    "        json.dump(labelme_data, f, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76c15055",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T09:48:07.213071Z",
     "start_time": "2024-03-15T09:48:07.208841Z"
    }
   },
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "\n",
    "# 设置文件夹路径和输出文件夹路径\n",
    "txt_folder = \"/home/wssysh/temp/lanedata/labels/\"  # 存放LabelImg标注的txt文件的文件夹路径\n",
    "output_folder = \"/home/wssysh/temp/lanedata/temp/\"  # 输出LabelMe标注的json文件的文件夹路径\n",
    "img_folder = \"/home/wssysh/temp/lanedata/imgs/\"  #存放对应标签的图片文件夹路径\n",
    "\n",
    "# 创建输出文件夹\n",
    "if not os.path.exists(output_folder):\n",
    "    os.makedirs(output_folder)\n",
    "\n",
    "# 遍历txt文件夹中的所有文件\n",
    "for filename in os.listdir(txt_folder):\n",
    "    if filename.endswith('.txt'):\n",
    "        # 生成对应的输出文件名\n",
    "        output_filename = os.path.splitext(filename)[0] + '.json'\n",
    "\n",
    "        # 读取txt文件\n",
    "        txt_file = os.path.join(txt_folder, filename)\n",
    "        data = read_txt_file(txt_file)\n",
    "\n",
    "        # 设置图片路径和尺寸\n",
    "        image_filename = os.path.splitext(\n",
    "            filename)[0] + '.jpg'  # 图片文件名与txt文件名相同，后缀为.jpg\n",
    "        image_path = os.path.join(img_folder, image_filename)\n",
    "        #  mage_size = (1440, 1080)  #(1280, 720)  # 根据实际情况修改\n",
    "        #  PIL.Image  获取 图像宽高\n",
    "        image_size = Image.open(image_path).size\n",
    "        # 转化为LabelMe格式\n",
    "        labelme_data = convert_to_labelme(data, image_filename, image_size)\n",
    "\n",
    "        # 保存为LabelMe JSON文件\n",
    "        output_file = os.path.join(output_folder, output_filename)\n",
    "        save_labelme_json(labelme_data, image_path, output_file)\n",
    "print(\"OK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a8d32b8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T09:00:54.420880Z",
     "start_time": "2024-03-15T09:00:54.410023Z"
    }
   },
   "outputs": [],
   "source": [
    "# 测试片段\n",
    "from PIL import Image\n",
    "img_path = \"/home/wssysh/Data/toSucceed/f0027.jpg\"\n",
    "img = Image.open(img_path)\n",
    "\n",
    "print(img.size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8514b0c7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T09:06:24.342942Z",
     "start_time": "2024-03-15T09:06:24.340542Z"
    }
   },
   "outputs": [],
   "source": [
    "# 测试片段\n",
    "i_num = 0\n",
    "for ind_num in range(4,47,3):\n",
    "    i_num=i_num + 1\n",
    "    print(ind_num,i_num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b578f8eb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T09:14:19.476125Z",
     "start_time": "2024-03-15T09:14:19.473875Z"
    }
   },
   "outputs": [],
   "source": [
    "# 测试片段\n",
    "labels=[\"l1\",\"l2\",\"l3\",\"l4\"]\n",
    "print(labels[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef2fe0a5",
   "metadata": {},
   "source": [
    "#### labelme json to YOLO txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33aad1a8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:10:43.727509Z",
     "start_time": "2024-03-15T10:10:43.724854Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import shutil\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "Dataset_root = '/home/wssysh/temp/lanedata/jsons/'\n",
    "os.chdir(Dataset_root)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a31bd30",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:10:46.346125Z",
     "start_time": "2024-03-15T10:10:46.343851Z"
    }
   },
   "outputs": [],
   "source": [
    "# 框的类别\n",
    "bbox_class = {\n",
    "    'roi':0  \n",
    "}\n",
    "# 关键点的类别\n",
    "keypoint_class = [\n",
    "    'llane0', 'llane1', 'llane2', 'llane3', 'llane4',\n",
    "    'lane0', 'lane1','lane2', 'lane3', 'lane4',\n",
    "    'rlane0', 'rlane1', 'rlane2', 'rlane3', 'rlane4'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "404202d7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:07:25.976847Z",
     "start_time": "2024-03-15T10:07:25.966118Z"
    }
   },
   "outputs": [],
   "source": [
    "os.mkdir('labels')\n",
    "os.mkdir('labels/train')\n",
    "os.mkdir('labels/val')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6468e978",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:10:50.899120Z",
     "start_time": "2024-03-15T10:10:50.783222Z"
    }
   },
   "outputs": [],
   "source": [
    "!ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d16cecc9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:10:54.504477Z",
     "start_time": "2024-03-15T10:10:54.496961Z"
    },
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def process_single_json(labelme_path, save_folder):\n",
    "    \n",
    "    with open(labelme_path, 'r', encoding='utf-8') as f:\n",
    "        labelme = json.load(f)\n",
    "\n",
    "    img_width = labelme['imageWidth']   # 图像宽度\n",
    "    img_height = labelme['imageHeight'] # 图像高度\n",
    "\n",
    "    # 生成 YOLO 格式的 txt 文件\n",
    "    suffix = labelme_path.split('.')[-2]\n",
    "    yolo_txt_path = suffix + '.txt'\n",
    "\n",
    "    with open(yolo_txt_path, 'w', encoding='utf-8') as f:\n",
    "\n",
    "        for each_ann in labelme['shapes']: # 遍历每个标注\n",
    "\n",
    "            if each_ann['shape_type'] == 'rectangle': # 每个框，在 txt 里写一行\n",
    "\n",
    "                yolo_str = ''\n",
    "\n",
    "                ## 框的信息\n",
    "                # 框的类别 ID\n",
    "                bbox_class_id = bbox_class[each_ann['label']]\n",
    "                yolo_str += '{} '.format(bbox_class_id)\n",
    "                # 左上角和右下角的 XY 像素坐标\n",
    "                bbox_top_left_x = int(min(each_ann['points'][0][0], each_ann['points'][1][0]))\n",
    "                bbox_bottom_right_x = int(max(each_ann['points'][0][0], each_ann['points'][1][0]))\n",
    "                bbox_top_left_y = int(min(each_ann['points'][0][1], each_ann['points'][1][1]))\n",
    "                bbox_bottom_right_y = int(max(each_ann['points'][0][1], each_ann['points'][1][1]))\n",
    "                # 框中心点的 XY 像素坐标\n",
    "                bbox_center_x = int((bbox_top_left_x + bbox_bottom_right_x) / 2)\n",
    "                bbox_center_y = int((bbox_top_left_y + bbox_bottom_right_y) / 2)\n",
    "                # 框宽度\n",
    "                bbox_width = bbox_bottom_right_x - bbox_top_left_x\n",
    "                # 框高度\n",
    "                bbox_height = bbox_bottom_right_y - bbox_top_left_y\n",
    "                # 框中心点归一化坐标\n",
    "                bbox_center_x_norm = bbox_center_x / img_width\n",
    "                bbox_center_y_norm = bbox_center_y / img_height\n",
    "                # 框归一化宽度\n",
    "                bbox_width_norm = bbox_width / img_width\n",
    "                # 框归一化高度\n",
    "                bbox_height_norm = bbox_height / img_height\n",
    "\n",
    "                yolo_str += '{:.5f} {:.5f} {:.5f} {:.5f} '.format(bbox_center_x_norm, bbox_center_y_norm, bbox_width_norm, bbox_height_norm)\n",
    "\n",
    "                ## 找到该框中所有关键点，存在字典 bbox_keypoints_dict 中\n",
    "                bbox_keypoints_dict = {}\n",
    "                for each_ann in labelme['shapes']: # 遍历所有标注\n",
    "                    if each_ann['shape_type'] == 'point': # 筛选出关键点标注\n",
    "                        # 关键点XY坐标、类别\n",
    "                        x = int(each_ann['points'][0][0])\n",
    "                        y = int(each_ann['points'][0][1])\n",
    "                        label = each_ann['label']\n",
    "                        if (x>bbox_top_left_x) & (x<bbox_bottom_right_x) & (y<bbox_bottom_right_y) & (y>bbox_top_left_y): # 筛选出在该个体框中的关键点\n",
    "                            bbox_keypoints_dict[label] = [x, y]\n",
    "\n",
    "                ## 把关键点按顺序排好\n",
    "                for each_class in keypoint_class: # 遍历每一类关键点\n",
    "                    if each_class in bbox_keypoints_dict:\n",
    "                        keypoint_x_norm = bbox_keypoints_dict[each_class][0] / img_width\n",
    "                        keypoint_y_norm = bbox_keypoints_dict[each_class][1] / img_height\n",
    "                        yolo_str += '{:.5f} {:.5f} {} '.format(keypoint_x_norm, keypoint_y_norm, 2) # 2-可见不遮挡 1-遮挡 0-没有点\n",
    "                    else: # 不存在的点，一律为0\n",
    "                        yolo_str += '0 0 0 '\n",
    "                # 写入 txt 文件中\n",
    "                f.write(yolo_str + '\\n')\n",
    "                \n",
    "    shutil.move(yolo_txt_path, save_folder)\n",
    "    print('{} --> {} 转换完成'.format(labelme_path, yolo_txt_path))    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b55b16c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T10:10:56.624791Z",
     "start_time": "2024-03-15T10:10:56.615743Z"
    }
   },
   "outputs": [],
   "source": [
    "# 部分由于预打标检出失败，需提前进行筛选\n",
    "# 通过文件大小排序即可快速清理\n",
    "save_folder = '/home/wssysh/temp/lanedata/labels/'\n",
    "for labelme_path in os.listdir():\n",
    "    try:\n",
    "        process_single_json(labelme_path, save_folder)\n",
    "    except:\n",
    "        print('******有误******', labelme_path)\n",
    "print('YOLO格式的txt标注文件已保存至 ', save_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99cdc395",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "625acb89",
   "metadata": {},
   "source": [
    "### 划分训练集和验证集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5575ef6d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:29:02.879278Z",
     "start_time": "2024-03-15T05:29:02.874618Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import random\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5f66e8f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:29:20.876034Z",
     "start_time": "2024-03-15T05:29:20.874183Z"
    }
   },
   "outputs": [],
   "source": [
    "Dataset_root = '/home/wssysh/temp/lanedata/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7cc0209",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:29:26.878503Z",
     "start_time": "2024-03-15T05:29:26.875100Z"
    }
   },
   "outputs": [],
   "source": [
    "os.chdir(os.path.join(Dataset_root, 'imgs'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53a3d75f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:29:32.138042Z",
     "start_time": "2024-03-15T05:29:32.123759Z"
    }
   },
   "outputs": [],
   "source": [
    "os.listdir()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18721aae",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:30:47.236595Z",
     "start_time": "2024-03-15T05:30:47.234628Z"
    }
   },
   "outputs": [],
   "source": [
    "# 划分\n",
    "test_frac = 0.1  # 测试集比例\n",
    "random.seed(123) # 随机数种子，便于复现"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5aec9720",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:30:49.344836Z",
     "start_time": "2024-03-15T05:30:49.343061Z"
    }
   },
   "outputs": [],
   "source": [
    "folder = '.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e33b044",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:30:51.906197Z",
     "start_time": "2024-03-15T05:30:51.902969Z"
    }
   },
   "outputs": [],
   "source": [
    "img_paths = os.listdir(folder)\n",
    "random.shuffle(img_paths) # 随机打乱\n",
    "\n",
    "val_number = int(len(img_paths) * test_frac) # 测试集文件个数\n",
    "train_files = img_paths[val_number:]         # 训练集文件名列表\n",
    "val_files = img_paths[:val_number]           # 测试集文件名列表\n",
    "\n",
    "print('数据集文件总数', len(img_paths))\n",
    "print('训练集文件个数', len(train_files))\n",
    "print('测试集文件个数', len(val_files))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e0d3245",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:31:14.901862Z",
     "start_time": "2024-03-15T05:31:14.890574Z"
    }
   },
   "outputs": [],
   "source": [
    "os.mkdir('train')\n",
    "for each in tqdm(train_files):\n",
    "    shutil.move(each, 'train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94ea7274",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:31:27.706750Z",
     "start_time": "2024-03-15T05:31:27.702300Z"
    }
   },
   "outputs": [],
   "source": [
    "os.mkdir('val')\n",
    "for each in tqdm(val_files):\n",
    "    shutil.move(each, 'val')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52ae896a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:32:32.998560Z",
     "start_time": "2024-03-15T05:32:32.996699Z"
    }
   },
   "outputs": [],
   "source": [
    "# 对标签进行处理\n",
    "os.chdir('../labels/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b248b00",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:33:28.545376Z",
     "start_time": "2024-03-15T05:33:28.536738Z"
    }
   },
   "outputs": [],
   "source": [
    "os.mkdir('train')\n",
    "for each in tqdm(train_files):\n",
    "    srt_path = each.split('.')[0] + '.txt'\n",
    "    shutil.move(srt_path, 'train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbc2430a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:33:41.642240Z",
     "start_time": "2024-03-15T05:33:41.637575Z"
    }
   },
   "outputs": [],
   "source": [
    "os.mkdir('val')\n",
    "for each in tqdm(val_files):\n",
    "    srt_path = each.split('.')[0] + '.txt'\n",
    "    shutil.move(srt_path, 'val')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "963abd86",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T05:34:03.738954Z",
     "start_time": "2024-03-15T05:34:03.626972Z"
    }
   },
   "outputs": [],
   "source": [
    "os.chdir('../')\n",
    "!pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de71ea09",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7ddcabf",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3fc1fe9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "7841443c",
   "metadata": {},
   "source": [
    "## 文件重命名"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82033bf8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.153526Z",
     "start_time": "2022-07-16T08:19:43.153509Z"
    }
   },
   "outputs": [],
   "source": [
    "# 利用linux命令对图片重命名\n",
    "import os\n",
    "root_path = \"/home/wssysh/Downloads/temp1/\"\n",
    "# 文件夹计数器\n",
    "count = 141\n",
    "\n",
    "for dir in os.listdir(root_path):\n",
    "    #print(dir)\n",
    "    count += 1\n",
    "    imgpath = os.path.join(root_path, dir)\n",
    "    print(\"cd {}&&for FILENAME in *; do mv $FILENAME {:0>2d}_$FILENAME; done\".format(imgpath,count))\n",
    "\n",
    "    os.system(\"cd {}&&for FILENAME in *;do mv $FILENAME {:0>2d}_$FILENAME;done\".format(imgpath, count))\n",
    "    \n",
    "print(\"OK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9627f12",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.154789Z",
     "start_time": "2022-07-16T08:19:43.154773Z"
    }
   },
   "outputs": [],
   "source": [
    "# $ mv ./*/*jpg /home/wssysh/data/datasets/culine/\n",
    "540/1.5 # h ->352"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fec4fe5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.156118Z",
     "start_time": "2022-07-16T08:19:43.156102Z"
    }
   },
   "outputs": [],
   "source": [
    "# 图片 resize\n",
    "import os\n",
    "imgpath = \"/home/wssysh/Downloads/temp/\"\n",
    "for img in os.listdir(imgpath):\n",
    "    os.system(\"cd {0}&& convert -resize 960x512! {1} /home/wssysh/Downloads/temp1/{1}\".format(imgpath,img))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25d5f0f4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.157373Z",
     "start_time": "2022-07-16T08:19:43.157357Z"
    }
   },
   "outputs": [],
   "source": [
    "# 注意顺序，另外在实验中发现不能成功调用yolo.detectimage()\n",
    "def xyxy2xywh(size, box):\n",
    "    dw = 1./size[0]\n",
    "    dh = 1./size[1]\n",
    "    x = (box[0] + box[1])/2.0\n",
    "    y = (box[2] + box[3])/2.0\n",
    "    w = box[1] - box[0]\n",
    "    h = box[3] - box[2]\n",
    "    x = x*dw\n",
    "    w = w*dw\n",
    "    y = y*dh\n",
    "    h = h*dh\n",
    "    return (x,y,w,h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49e9aa3c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.158924Z",
     "start_time": "2022-07-16T08:19:43.158907Z"
    }
   },
   "outputs": [],
   "source": [
    "# 移动 筛选后的反转图片\n",
    "import os\n",
    "import shutil\n",
    "# 路径后面带/\n",
    "img_path = \"/home/wssysh/data/videos/keras_yolo/bs_yolov5/dataset/VOCdevkit/VOC2007/JPEGImages/\"\n",
    "imgnames = os.listdir(img_path)\n",
    "for imgname in imgnames:\n",
    "    pic_path = img_path.replace(\"JPEGImages\", \"rpics\") + \"r\" + imgname\n",
    "    if os.path.isfile(pic_path):\n",
    "        shutil.move(pic_path, img_path.replace(\"JPEGImages\", \"ceui\"))\n",
    "    else:\n",
    "        print(imgname)\n",
    "print(\"OK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "818aff2f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.160018Z",
     "start_time": "2022-07-16T08:19:43.160008Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "# 路径后面带/\n",
    "img_path = \"/home/wssysh/data/videos/keras_yolo/bs_yolov5/dataset/VOCdevkit/VOC2007/Annotations/\"\n",
    "imgnames = os.listdir(img_path)\n",
    "for imgname in imgnames:\n",
    "    shutil.move(\n",
    "        img_path.replace(\"Annotations\", \"JPEGImages\") +\n",
    "        imgname.replace(\"xml\", \"jpg\"), img_path.replace(\"Annotations\", \"pics\"))\n",
    "print(\"OK\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4925a05",
   "metadata": {},
   "source": [
    "# 数据库\n",
    "\n",
    "- 添加数据库，存储试卷成绩"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61b8f798",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.160710Z",
     "start_time": "2022-07-16T08:19:43.160700Z"
    }
   },
   "outputs": [],
   "source": [
    "import sqlite3\n",
    "# 创建与数据库的连接\n",
    "conn = sqlite3.connect('test.db')\n",
    "#创建一个游标 cursor\n",
    "cur = conn.cursor()\n",
    "\n",
    "# 建表的sql语句\n",
    "sql_text_1 = '''CREATE TABLE scores\n",
    "           (姓名 TEXT,\n",
    "            班级 TEXT,\n",
    "            性别 TEXT,\n",
    "            语文 NUMBER,\n",
    "            数学 NUMBER,\n",
    "            英语 NUMBER);'''\n",
    "# 执行sql语句\n",
    "cur.execute(sql_text_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0633ff6c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.161505Z",
     "start_time": "2022-07-16T08:19:43.161494Z"
    }
   },
   "outputs": [],
   "source": [
    "# 插入单条数据\n",
    "sql_text_2 = \"INSERT INTO scores VALUES('A', '一班', '男', 96, 94, 98)\"\n",
    "cur.execute(sql_text_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cc4e542",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.162328Z",
     "start_time": "2022-07-16T08:19:43.162318Z"
    }
   },
   "outputs": [],
   "source": [
    "# 执行以下语句插入多条数据：\n",
    "\n",
    "data = [('B', '一班', '女', 78, 87, 85),\n",
    "        ('C', '一班', '男', 98, 84, 90),\n",
    "        ]\n",
    "cur.executemany('INSERT INTO scores VALUES (?,?,?,?,?,?)', data)\n",
    "# 连接完数据库并不会自动提交，所以需要手动 commit 你的改动conn.commit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d133c3f9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.163135Z",
     "start_time": "2022-07-16T08:19:43.163124Z"
    }
   },
   "outputs": [],
   "source": [
    "# 查询数学成绩大于90分的学生\n",
    "sql_text_3 = \"SELECT * FROM scores WHERE 数学>90\"\n",
    "cur.execute(sql_text_3)\n",
    "# 获取查询结果\n",
    "cur.fetchall()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d985fee",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.163774Z",
     "start_time": "2022-07-16T08:19:43.163764Z"
    }
   },
   "outputs": [],
   "source": [
    "# 对数据库做改动后（比如建表、插数等），都需要手动提交改动，否则无法将数据保存到数据库。\n",
    "\n",
    "# 提交改动的方法\n",
    "conn.commit()\n",
    "使用完数据库之后，需要关闭游标和连接：\n",
    "\n",
    "# 关闭游标\n",
    "cur.close()\n",
    "# 关闭连接\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc217abd",
   "metadata": {},
   "source": [
    "# 环境测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc7a7503",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.164563Z",
     "start_time": "2022-07-16T08:19:43.164553Z"
    }
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c21def3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-10-15T00:32:50.847169Z",
     "start_time": "2022-10-15T00:32:50.227299Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de0463a9",
   "metadata": {},
   "source": [
    "# 图片矫正与边缘提取"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8aecefd",
   "metadata": {},
   "source": [
    "## imgshow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df3d38bf",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-23T01:29:10.774630Z",
     "start_time": "2022-07-23T01:29:10.769906Z"
    }
   },
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "import cv2\n",
    "\n",
    "def imgshow(img, winname=\"test\", showt=0):\n",
    "    cv2.namedWindow(winname, cv2.WINDOW_NORMAL)\n",
    "    cv2.resizeWindow(winname, 1600, 1200)\n",
    "    cv2.imshow(winname, img)\n",
    "    cv2.waitKey(showt)\n",
    "    cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55fbaf3e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T09:12:04.712980Z",
     "start_time": "2022-07-16T09:12:02.077199Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import cv2\n",
    "from IPython.display import clear_output\n",
    "\n",
    "img = cv2.imread(\"/home/wssysh/data/to_learn/notebooks/output/no_crop/date21/IMG_20220621_182630.jpg\")\n",
    "\n",
    "if img is not None:\n",
    "    imgshow(img)\n",
    "\n",
    "clear_output(wait=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae52647d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.167779Z",
     "start_time": "2022-07-16T08:19:43.167769Z"
    }
   },
   "outputs": [],
   "source": [
    "cv2.imwrite(\"../output/crop\" + '.jpg',dst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90005172",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T09:16:05.358390Z",
     "start_time": "2022-07-16T09:15:56.474536Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "\n",
    "filepath = \"../output/no_crop/date21/\"\n",
    "\n",
    "# original pts 大致变换\n",
    "# date21\n",
    "# 变换前角点位置 变换之后的图上四个点的位置\n",
    "pts_o = np.float32([[1460, 0], [3270, 0], [0, cols * 5 / 6],[4620, cols * 5 / 6]])\n",
    "pts_d = np.float32([[0, 0], [1280, 0], [0, 1920], [1280, 1920]]) \n",
    "\n",
    "for file in os.listdir(filepath):\n",
    "    img = cv2.imread(os.path.join(filepath, file))\n",
    "    dst = cv2.warpPerspective(img, M, (1280, 1920))  # 输出dst的尺寸\n",
    "    cv2.imwrite(\"../output/crop21/\" + file + '.jpg', dst)\n",
    "print(\"OK\")    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dec1f01",
   "metadata": {},
   "source": [
    "## 透射变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50b27475",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-21T12:51:13.236307Z",
     "start_time": "2022-07-21T12:51:09.679705Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2 as cv\n",
    "import matplotlib.pyplot as plt\n",
    "# 1 读取图像\n",
    "img = cv.imread(\"../output/testdata0721/IMG_20220721_191722.jpg\")\n",
    "imgshow(img,\"ori_img\")\n",
    "\n",
    "clear_output(wait=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b04f436e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-21T01:12:01.006056Z",
     "start_time": "2022-07-21T01:11:59.586118Z"
    }
   },
   "outputs": [],
   "source": [
    "# 2 透射变换\n",
    "rows, cols = img.shape[:2]\n",
    "\n",
    "# 2.1 创建变换矩阵\n",
    "pts_o = np.float32([[1426, 231], [3134, 248], [334, 3209],[4250, 3213]])\n",
    "pts_d = np.float32([[0, 0], [1280, 0], [0, 1920], [1280, 1920]]) \n",
    "\n",
    "T = cv.getPerspectiveTransform(pts_o, pts_d)\n",
    "print(T)\n",
    "# 2.2 进行变换\n",
    "dst = cv.warpPerspective(img, T, (1280, 1920))\n",
    "\n",
    "# 3 图像显示\n",
    "fig, axes=plt.subplots(nrows=1, ncols=2, figsize=(10, 8), dpi=100)\n",
    "axes[0].imshow(img[:, :, ::-1])\n",
    "axes[0].set_title(\"ori\")\n",
    "axes[1].imshow(dst[:, :, ::-1])\n",
    "axes[1].set_title(\"wrapped\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a4da539",
   "metadata": {
    "ExecuteTime": {
     "start_time": "2022-07-19T07:54:54.195Z"
    }
   },
   "outputs": [],
   "source": [
    "imgshow(dst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8984824c",
   "metadata": {},
   "outputs": [],
   "source": [
    "[[ 1.15996013e+00  6.97541818e-01 -7.87597720e+02]\n",
    " [ 5.56212834e-02  2.28569345e+00 -2.17284632e+02]\n",
    " [ 1.50759522e-05  3.84436178e-04  1.00000000e+00]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28fe2005",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-21T01:12:34.881980Z",
     "start_time": "2022-07-21T01:12:34.877448Z"
    }
   },
   "outputs": [],
   "source": [
    "# 变换矩阵\n",
    "testT = T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "214cc1da",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-21T01:13:35.507006Z",
     "start_time": "2022-07-21T01:13:29.866687Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "\n",
    "filepath = \"../output/no_crop/date22/\"\n",
    "for file in os.listdir(filepath):\n",
    "    img = cv2.imread(os.path.join(filepath, file))\n",
    "    dst = cv2.warpPerspective(img, testT,(1280, 1920) )  # 输出dst的尺寸\n",
    "    cv2.imwrite(\"../output/crop22/\" + file + '.jpg', dst)\n",
    "print(\"OK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdd33f41",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-21T12:50:53.505526Z",
     "start_time": "2022-07-21T12:50:52.235619Z"
    }
   },
   "outputs": [],
   "source": [
    "#\n",
    "import PIL.Image as Image\n",
    "img = Image.open(\"../output/no_crop/date21/IMG_20220621_182612.jpg\")\n",
    "img.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9157104",
   "metadata": {},
   "source": [
    "## 霍夫变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a257aa21",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-23T01:35:28.708099Z",
     "start_time": "2022-07-23T01:29:14.591192Z"
    }
   },
   "outputs": [],
   "source": [
    "# coding=utf-8\n",
    "# 导入相应的python包\n",
    "import cv2 \n",
    "import numpy as np \n",
    "  \n",
    "# 读取输入图片\n",
    "img = cv2.imread('../output/testdata0721/IMG_20220721_191722.jpg')\n",
    "img_resize = cv2.resize(img,(800,600))\n",
    "imgshow(img_resize,\"img_resize\")\n",
    "\n",
    "ksize = (10, 10) \n",
    "# Using cv2.blur() method\n",
    "img_blur = cv2.blur(img_resize, ksize, cv2.BORDER_DEFAULT)\n",
    "# imgshow(img_blur,\"img_blur\")\n",
    "\n",
    "\n",
    "# 将彩色图片灰度化\n",
    "gray = cv2.cvtColor(img_blur,cv2.COLOR_BGR2GRAY) \n",
    "imgshow(gray,\"gray\")\n",
    "\n",
    "# _,img_bin =cv2.threshold(gray, 140, 255, cv2.THRESH_BINARY);\n",
    "img_bin = cv2.adaptiveThreshold(gray,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY,11,5)\n",
    "imgshow(img_bin,\"img_bin\")\n",
    "\n",
    "# 腐蚀膨胀操作需注意主体二值化，选择性使用\n",
    "# 膨胀\n",
    "kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))\n",
    "dst = cv2.dilate(img_bin, kernel)\n",
    "imgshow(dst,\"dilate_demo\")\n",
    "\n",
    "# 腐蚀\n",
    "kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))\n",
    "dst = cv2.erode(dst, kernel)\n",
    "imgshow(dst,\"erode_demo\")\n",
    "\n",
    "\n",
    "# 使用Canny边缘检测 \n",
    "edges = cv2.Canny(img_bin,50,200,apertureSize = 3) \n",
    "imgshow(edges,\"edges\")\n",
    "\n",
    "# 进行Hough_line直线检测\n",
    "lines = cv2.HoughLinesP(edges,1,np.pi/180, 80, 30, 10) \n",
    "\n",
    "# 遍历每一条直线\n",
    "for i in range(len(lines)): \n",
    "\tcv2.line(img,(lines[i, 0, 0],lines[i, 0, 1]), (lines[i, 0, 2],lines[i, 0, 3]), (0,255,0),2) \n",
    "\n",
    "imgshow(img)\n",
    "# 保存结果\n",
    "cv2.imwrite('test3_r.jpg', img)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a1c67f7",
   "metadata": {},
   "source": [
    "## LSD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4ee151d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-23T01:28:27.257188Z",
     "start_time": "2022-07-23T01:28:06.720692Z"
    }
   },
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "# 读取输入图片\n",
    "img = cv2.imread(\"../output/testdata0721/IMG_20220721_191722.jpg\")\n",
    "# 将彩色图片转换为灰度图片\n",
    "img_resize = cv2.resize(img,(800,600))\n",
    "\n",
    "img0 = cv2.blur(img_resize, (15, 15), 0)\n",
    "\n",
    "#img0 = cv2.medianBlur(img, 21)\n",
    "\n",
    "img = cv2.cvtColor(img0,cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "#img_bin = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, blockSize=255, C=90)\n",
    "\n",
    "# 创建一个LSD对象\n",
    "lsd = cv2.createLineSegmentDetector(0)\n",
    "# 执行检测结果\n",
    "#dlines = lsd.detect(img_bin)\n",
    "dlines = lsd.detect(img)\n",
    "# 绘制检测结果\n",
    "for dline in dlines[0]:\n",
    "    x0 = int(round(dline[0][0]))\n",
    "    y0 = int(round(dline[0][1]))\n",
    "    x1 = int(round(dline[0][2]))\n",
    "    y1 = int(round(dline[0][3]))\n",
    "    cv2.line(img0, (x0, y0), (x1,y1), (0,255,0), 1, cv2.LINE_AA)\n",
    "\n",
    "# 显示并保存结果\n",
    "#cv2.imwrite('3.JPG', img0)\n",
    "cv2.imshow(\"LSD\", img0)\n",
    "cv2.waitKey(0)\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "066f4f75",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.153526Z",
     "start_time": "2022-07-16T08:19:43.153509Z"
    }
   },
   "outputs": [],
   "source": [
    "# 利用linux命令对图片重命名\n",
    "import os\n",
    "root_path = \"/home/wssysh/Downloads/temp1/\"\n",
    "# 文件夹计数器\n",
    "count = 141\n",
    "\n",
    "for dir in os.listdir(root_path):\n",
    "    #print(dir)\n",
    "    count += 1\n",
    "    imgpath = os.path.join(root_path, dir)\n",
    "    print(\"cd {}&&for FILENAME in *; do mv $FILENAME {:0>2d}_$FILENAME; done\".format(imgpath,count))\n",
    "\n",
    "    os.system(\"cd {}&&for FILENAME in *;do mv $FILENAME {:0>2d}_$FILENAME;done\".format(imgpath, count))\n",
    "    \n",
    "print(\"OK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d3eac69",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.154789Z",
     "start_time": "2022-07-16T08:19:43.154773Z"
    }
   },
   "outputs": [],
   "source": [
    "# $ mv ./*/*jpg /home/wssysh/data/datasets/culine/\n",
    "540/1.5 # h ->352"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd28fa02",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.156118Z",
     "start_time": "2022-07-16T08:19:43.156102Z"
    }
   },
   "outputs": [],
   "source": [
    "# 图片 resize\n",
    "import os\n",
    "imgpath = \"/home/wssysh/Downloads/temp/\"\n",
    "for img in os.listdir(imgpath):\n",
    "    os.system(\"cd {0}&& convert -resize 960x512! {1} /home/wssysh/Downloads/temp1/{1}\".format(imgpath,img))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9dd10b1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.157373Z",
     "start_time": "2022-07-16T08:19:43.157357Z"
    }
   },
   "outputs": [],
   "source": [
    "# 注意顺序，另外在实验中发现不能成功调用yolo.detectimage()\n",
    "def xyxy2xywh(size, box):\n",
    "    dw = 1./size[0]\n",
    "    dh = 1./size[1]\n",
    "    x = (box[0] + box[1])/2.0\n",
    "    y = (box[2] + box[3])/2.0\n",
    "    w = box[1] - box[0]\n",
    "    h = box[3] - box[2]\n",
    "    x = x*dw\n",
    "    w = w*dw\n",
    "    y = y*dh\n",
    "    h = h*dh\n",
    "    return (x,y,w,h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "457ffe7b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.158924Z",
     "start_time": "2022-07-16T08:19:43.158907Z"
    }
   },
   "outputs": [],
   "source": [
    "# 移动 筛选后的反转图片\n",
    "import os\n",
    "import shutil\n",
    "# 路径后面带/\n",
    "img_path = \"/home/wssysh/data/videos/keras_yolo/bs_yolov5/dataset/VOCdevkit/VOC2007/JPEGImages/\"\n",
    "imgnames = os.listdir(img_path)\n",
    "for imgname in imgnames:\n",
    "    pic_path = img_path.replace(\"JPEGImages\", \"rpics\") + \"r\" + imgname\n",
    "    if os.path.isfile(pic_path):\n",
    "        shutil.move(pic_path, img_path.replace(\"JPEGImages\", \"ceui\"))\n",
    "    else:\n",
    "        print(imgname)\n",
    "print(\"OK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3af0e347",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.160018Z",
     "start_time": "2022-07-16T08:19:43.160008Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "# 路径后面带/\n",
    "img_path = \"/home/wssysh/data/videos/keras_yolo/bs_yolov5/dataset/VOCdevkit/VOC2007/Annotations/\"\n",
    "imgnames = os.listdir(img_path)\n",
    "for imgname in imgnames:\n",
    "    shutil.move(\n",
    "        img_path.replace(\"Annotations\", \"JPEGImages\") +\n",
    "        imgname.replace(\"xml\", \"jpg\"), img_path.replace(\"Annotations\", \"pics\"))\n",
    "print(\"OK\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40ad0201",
   "metadata": {},
   "source": [
    "# 数据库\n",
    "\n",
    "- 添加数据库，存储试卷成绩"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60e53822",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-18T04:52:39.524771Z",
     "start_time": "2023-09-18T04:52:39.419252Z"
    }
   },
   "outputs": [],
   "source": [
    "import sqlite3\n",
    "# 创建与数据库的连接\n",
    "conn = sqlite3.connect('test.db')\n",
    "#创建一个游标 cursor\n",
    "cur = conn.cursor()\n",
    "\n",
    "# 建表的sql语句\n",
    "sql_text_1 = '''CREATE TABLE scores\n",
    "           (姓名 TEXT,\n",
    "            班级 TEXT,\n",
    "            性别 TEXT,\n",
    "            语文 NUMBER,\n",
    "            数学 NUMBER,\n",
    "            英语 NUMBER);'''\n",
    "# 执行sql语句\n",
    "cur.execute(sql_text_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3489ab51",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-18T04:52:42.650796Z",
     "start_time": "2023-09-18T04:52:42.647102Z"
    }
   },
   "outputs": [],
   "source": [
    "# 插入单条数据\n",
    "sql_text_2 = \"INSERT INTO scores VALUES('A', '一班', '男', 96, 94, 98)\"\n",
    "cur.execute(sql_text_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfdd48df",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-18T04:52:46.354334Z",
     "start_time": "2023-09-18T04:52:46.350562Z"
    }
   },
   "outputs": [],
   "source": [
    "# 执行以下语句插入多条数据：\n",
    "\n",
    "data = [('B', '一班', '女', 78, 87, 85),\n",
    "        ('C', '一班', '男', 98, 84, 90),\n",
    "        ]\n",
    "cur.executemany('INSERT INTO scores VALUES (?,?,?,?,?,?)', data)\n",
    "# 连接完数据库并不会自动提交，所以需要手动 commit 你的改动conn.commit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0811f64",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-18T04:52:53.127289Z",
     "start_time": "2023-09-18T04:52:53.118429Z"
    }
   },
   "outputs": [],
   "source": [
    "# 查询数学成绩大于90分的学生\n",
    "sql_text_3 = \"SELECT * FROM scores WHERE 数学>90\"\n",
    "cur.execute(sql_text_3)\n",
    "# 获取查询结果\n",
    "cur.fetchall()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6436f42",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-18T04:53:08.091316Z",
     "start_time": "2023-09-18T04:53:08.027580Z"
    }
   },
   "outputs": [],
   "source": [
    "# 对数据库做改动后（比如建表、插数等），都需要手动提交改动，否则无法将数据保存到数据库。\n",
    "\n",
    "# 提交改动的方法\n",
    "conn.commit()\n",
    "# 使用完数据库之后，需要关闭游标和连接：\n",
    "\n",
    "# 关闭游标\n",
    "cur.close()\n",
    "# 关闭连接\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b7484d2",
   "metadata": {},
   "source": [
    "# 图片矫正与边缘提取"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3c43124",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:58:04.005850Z",
     "start_time": "2022-07-16T08:58:03.995999Z"
    }
   },
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "import cv2\n",
    "\n",
    "def imgshow(img, winname=\"test\", showt=0):\n",
    "    cv2.namedWindow(winname, cv2.WINDOW_NORMAL)\n",
    "    cv2.resizeWindow(winname, 1600, 1200)\n",
    "    cv2.imshow(winname, img)\n",
    "    cv2.waitKey(showt)\n",
    "    cv2.destroyAllWindows()\n",
    "    clear_output(wait=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc6960d4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T09:12:04.712980Z",
     "start_time": "2022-07-16T09:12:02.077199Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import cv2\n",
    "from IPython.display import clear_output\n",
    "\n",
    "img = cv2.imread(\"/home/wssysh/data/to_learn/notebooks/output/no_crop/date21/IMG_20220621_182630.jpg\")\n",
    "\n",
    "if img is not None:\n",
    "    imgshow(img)\n",
    "\n",
    "clear_output(wait=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34b9716b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T09:12:29.660824Z",
     "start_time": "2022-07-16T09:12:21.145176Z"
    }
   },
   "outputs": [],
   "source": [
    "rows, cols = img.shape[:2]\n",
    "\n",
    "# original pts\n",
    "pts_o = np.float32([[1460, 0], [3270, 0], [0, cols*5/6], [4620, cols*5/6]])\n",
    "pts_d = np.float32([[0, 0], [1280, 0], [0, 1920], [1280, 1920]]) # 这是变换之后的图上四个点的位置\n",
    "\n",
    "# get transform matrix\n",
    "M = cv2.getPerspectiveTransform(pts_o, pts_d)\n",
    "# apply transformation\n",
    "dst = cv2.warpPerspective(img, M, (1280, 1920)) # 输出dst的尺寸\n",
    "\n",
    "imgshow(dst,\"dst\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ab62c48",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T08:19:43.167779Z",
     "start_time": "2022-07-16T08:19:43.167769Z"
    }
   },
   "outputs": [],
   "source": [
    "cv2.imwrite(\"../output/crop\" + '.jpg',dst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91fd8651",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T09:16:05.358390Z",
     "start_time": "2022-07-16T09:15:56.474536Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "\n",
    "filepath = \"../output/no_crop/date21/\"\n",
    "\n",
    "# original pts 大致变换\n",
    "# date21\n",
    "# 变换前角点位置 变换之后的图上四个点的位置\n",
    "pts_o = np.float32([[1460, 0], [3270, 0], [0, cols * 5 / 6],[4620, cols * 5 / 6]])\n",
    "pts_d = np.float32([[0, 0], [1280, 0], [0, 1920], [1280, 1920]]) \n",
    "\n",
    "for file in os.listdir(filepath):\n",
    "    img = cv2.imread(os.path.join(filepath, file))\n",
    "    dst = cv2.warpPerspective(img, M, (1280, 1920))  # 输出dst的尺寸\n",
    "    cv2.imwrite(\"../output/crop21/\" + file + '.jpg', dst)\n",
    "print(\"OK\")    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91d92a6b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-07-16T09:14:37.504668Z",
     "start_time": "2022-07-16T09:14:37.388701Z"
    }
   },
   "source": [
    "[sharp](https://zhuanlan.zhihu.com/p/347053207)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d903956c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T13:50:38.121078Z",
     "start_time": "2023-03-20T13:50:37.580032Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from skimage.io import imshow, imread\n",
    "from skimage.color import rgb2yuv, rgb2hsv, rgb2gray, yuv2rgb, hsv2rgb\n",
    "from scipy.signal import convolve2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74263881",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T13:50:46.324617Z",
     "start_time": "2023-03-20T13:50:46.139682Z"
    }
   },
   "outputs": [],
   "source": [
    "dog = imread('/home/wssysh/data/datasets/myvalid/1/1.jpeg')\n",
    "plt.figure(num=None, figsize=(8, 6), dpi=20)\n",
    "imshow(dog)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "133f4fbc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T13:51:05.900155Z",
     "start_time": "2023-03-20T13:51:05.897107Z"
    }
   },
   "outputs": [],
   "source": [
    "# 锐化\n",
    "sharpen = np.array([[0, -1, 0],\n",
    "                    [-1, 5, -1],\n",
    "                    [0, -1, 0]])\n",
    "                    \n",
    "# 高斯模糊\n",
    "gaussian = (1 / 16.0) * np.array([[1., 2., 1.],\n",
    "                                  [2., 4., 2.],\n",
    "                                  [1., 2., 1.]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "746e8905",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T14:03:03.214220Z",
     "start_time": "2023-03-20T14:03:03.164186Z"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,2, figsize = (7,5))\n",
    "ax[0].imshow(sharpen, cmap='gray')\n",
    "ax[0].set_title(f'Sharpen', fontsize = 18)\n",
    "    \n",
    "ax[1].imshow(gaussian, cmap='gray')\n",
    "ax[1].set_title(f'Gaussian Blur', fontsize = 18)\n",
    "    \n",
    "[axi.set_axis_off() for axi in ax.ravel()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d6a631",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T13:51:43.679593Z",
     "start_time": "2023-03-20T13:51:43.677547Z"
    }
   },
   "outputs": [],
   "source": [
    "def multi_convolver(image, kernel, iterations):\n",
    "    for i in range(iterations):\n",
    "        image = convolve2d(image, kernel, 'same', boundary = 'fill',\n",
    "                           fillvalue = 0)\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49c4b223",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T13:52:07.182708Z",
     "start_time": "2023-03-20T13:52:06.932382Z"
    }
   },
   "outputs": [],
   "source": [
    "multi_convolver(dog, gaussian, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "188d0a28",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T14:03:18.025129Z",
     "start_time": "2023-03-20T14:03:17.911389Z"
    }
   },
   "outputs": [],
   "source": [
    "dog_grey = rgb2gray(dog)\n",
    "plt.figure(num=None, figsize=(8, 6), dpi=27)\n",
    "imshow(dog_grey);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3eaf4857",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T14:03:31.605767Z",
     "start_time": "2023-03-20T14:03:31.427427Z"
    }
   },
   "outputs": [],
   "source": [
    "convolved_image = multi_convolver(dog_grey, gaussian, 5)\n",
    "plt.figure(num=None, figsize=(8, 6), dpi=27)\n",
    "imshow(convolved_image);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "941b6de5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T14:04:05.206174Z",
     "start_time": "2023-03-20T14:04:05.203438Z"
    }
   },
   "outputs": [],
   "source": [
    "def convolution_plotter(image, kernel):\n",
    "    iterations = [1,10,20,30]\n",
    "    f_size = 20\n",
    "    \n",
    "    fig, ax = plt.subplots(1,4, figsize = (7,5))\n",
    "    for n, ax in enumerate(ax.flatten()):\n",
    "        ax.set_title(f'Iteration : {iterations[n]}', fontsize =\n",
    "                     f_size)\n",
    "        ax.imshow(multi_convolver(image, kernel, iterations[n]), \n",
    "                  cmap='gray')\n",
    "        ax.set_axis_off()\n",
    "\n",
    "    fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ec8c966",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T13:53:45.979799Z",
     "start_time": "2023-03-20T13:53:44.705256Z"
    }
   },
   "outputs": [],
   "source": [
    "convolution_plotter(dog_grey, gaussian)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03460174",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T14:04:17.169687Z",
     "start_time": "2023-03-20T14:04:17.165366Z"
    }
   },
   "outputs": [],
   "source": [
    "def convolver_rgb(image, kernel, iterations = 1):\n",
    "    convolved_image_r = multi_convolver(image[:,:,0], kernel,\n",
    "                                        iterations)\n",
    "    convolved_image_g = multi_convolver(image[:,:,1], kernel, \n",
    "                                        iterations)\n",
    "    convolved_image_b  = multi_convolver(image[:,:,2], kernel, \n",
    "                                         iterations)\n",
    "    \n",
    "    reformed_image = np.dstack((np.rint(abs(convolved_image_r)), \n",
    "                                np.rint(abs(convolved_image_g)), \n",
    "                                np.rint(abs(convolved_image_b)))) / 255\n",
    "   \n",
    "                               \n",
    "                     \n",
    "    fig, ax = plt.subplots(1,3, figsize = (7,5))\n",
    "    \n",
    "    ax[0].imshow(abs(convolved_image_r), cmap='Reds')\n",
    "    ax[0].set_title(f'Red', fontsize = 15)\n",
    "    \n",
    "    ax[1].imshow(abs(convolved_image_g), cmap='Greens')\n",
    "    ax[1].set_title(f'Green', fontsize = 15)\n",
    "    \n",
    "    ax[2].imshow(abs(convolved_image_b), cmap='Blues')\n",
    "    ax[2].set_title(f'Blue', fontsize = 15)\n",
    "    \n",
    "    [axi.set_axis_off() for axi in ax.ravel()]\n",
    "    \n",
    "    return np.array(reformed_image).astype(np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70d70495",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T14:04:20.392490Z",
     "start_time": "2023-03-20T14:04:20.140053Z"
    }
   },
   "outputs": [],
   "source": [
    "convolved_rgb_sharpen = convolver_rgb(dog, sharpen, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39924971",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T14:04:32.750996Z",
     "start_time": "2023-03-20T14:04:32.747834Z"
    }
   },
   "outputs": [],
   "source": [
    "def convolver_rgb(image, kernel, iterations = 1):\n",
    "    img_yuv = rgb2yuv(image)   \n",
    "    img_yuv[:,:,0] = multi_convolver(img_yuv[:,:,0], kernel, \n",
    "                                     iterations)\n",
    "    final_image = yuv2rgb(img_yuv)\n",
    "                               \n",
    "    fig, ax = plt.subplots(1,2, figsize = (7,5))\n",
    "    \n",
    "    ax[0].imshow(image)\n",
    "    ax[0].set_title(f'Original', fontsize = 20)\n",
    "    \n",
    "    ax[1].imshow(final_image);\n",
    "    ax[1].set_title(f'YUV Adjusted, Iterations = {iterations}', \n",
    "                    fontsize = 20)\n",
    "    \n",
    "    [axi.set_axis_off() for axi in ax.ravel()]\n",
    "    \n",
    "    fig.tight_layout()\n",
    "    \n",
    "    return final_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9134c71",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T14:04:52.975104Z",
     "start_time": "2023-03-20T14:04:52.739198Z"
    }
   },
   "outputs": [],
   "source": [
    "final_image = convolver_rgb(dog, sharpen, iterations = 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1557ab0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T14:05:16.747646Z",
     "start_time": "2023-03-20T14:05:16.740027Z"
    }
   },
   "outputs": [],
   "source": [
    "def convolver_comparison(image, kernel, iterations = 1):\n",
    "    img_yuv = rgb2yuv(image)   \n",
    "    img_yuv[:,:,0] = multi_convolver(img_yuv[:,:,0], kernel, \n",
    "                      iterations)\n",
    "    final_image_yuv = yuv2rgb(img_yuv)\n",
    "    \n",
    "    img_hsv = rgb2hsv(image)   \n",
    "    img_hsv[:,:,2] = multi_convolver(img_hsv[:,:,2], kernel, \n",
    "                      iterations)\n",
    "    final_image_hsv = hsv2rgb(img_hsv)\n",
    "    \n",
    "    \n",
    "    \n",
    "    convolved_image_r = multi_convolver(image[:,:,0], kernel, \n",
    "                         iterations)\n",
    "    convolved_image_g = multi_convolver(image[:,:,1], kernel, \n",
    "                         iterations)\n",
    "    convolved_image_b  = multi_convolver(image[:,:,2], kernel,\n",
    "                         iterations)\n",
    "    \n",
    "    final_image_rgb = np.dstack((np.rint(abs(convolved_image_r)), \n",
    "                                 np.rint(abs(convolved_image_g)), \n",
    "                                 np.rint(abs(convolved_image_b))))/255\n",
    "                               \n",
    "    fig, ax = plt.subplots(2,2, figsize = (12,12))\n",
    "    \n",
    "    ax[0][0].imshow(image)\n",
    "    ax[0][0].set_title(f'Original', fontsize = 30)\n",
    "    \n",
    "    ax[0][1].imshow(final_image_rgb);\n",
    "    ax[0][1].set_title(f'RGB Adjusted, Iterations = {iterations}', \n",
    "                       fontsize = 30)\n",
    "    fig.tight_layout()\n",
    "    \n",
    "    ax[1][0].imshow(final_image_yuv)\n",
    "    ax[1][0].set_title(f'YUV Adjusted, Iterations = {iterations}', \n",
    "                       fontsize = 30)\n",
    "    \n",
    "    ax[1][1].imshow(final_image_hsv)\n",
    "    ax[1][1].set_title(f'HSV Adjusted, Iterations = {iterations}', \n",
    "                       fontsize = 30)\n",
    "    \n",
    "    [axi.set_axis_off() for axi in ax.ravel()]\n",
    "    \n",
    "    fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68d8429c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-03-20T14:05:23.072580Z",
     "start_time": "2023-03-20T14:05:18.991966Z"
    }
   },
   "outputs": [],
   "source": [
    "convolver_comparison(dog, sharpen, iterations = 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fb9876e",
   "metadata": {},
   "source": [
    "# 测试网络"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10697df2-e2e9-4b4e-9d95-9486b0a6976b",
   "metadata": {},
   "source": [
    "```sh\n",
    "#！/bin/bash\n",
    "# 关键点检测 训练脚本\n",
    "# 激活conda pytorch 环境\n",
    "conda activate pytorch \n",
    "# 调用YOLOv8 进行训练\n",
    "yolo pose train data=./lane.yaml model=yolov8n-pose.pt project=lanedet device=0 pretrained=True epochs=300 batch=8 device=0 close_mosaic=1\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86e28a5a",
   "metadata": {},
   "source": [
    "## Alexnet"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "206.8px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "oldHeight": 389.188,
   "position": {
    "height": "411.175px",
    "left": "1202.08px",
    "right": "20px",
    "top": "231.987px",
    "width": "800px"
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "varInspector_section_display": "block",
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
